nianzs
106 lines
3.9 KiB
Go
106 lines
3.9 KiB
Go
package kiro
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestRefreshSocialTokenInvalidGrantReturnsTypedError(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/refreshToken", r.URL.Path)
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
_, _ = w.Write([]byte(`{"error":"invalid_grant","message":"Invalid refresh token provided"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
previous := socialAuthEndpointURL
|
|
socialAuthEndpointURL = server.URL
|
|
t.Cleanup(func() { socialAuthEndpointURL = previous })
|
|
|
|
_, err := RefreshSocialToken(context.Background(), "", "revoked-refresh-token", "Google")
|
|
require.Error(t, err)
|
|
|
|
var invalid *RefreshTokenInvalidError
|
|
require.True(t, errors.As(err, &invalid))
|
|
require.Equal(t, http.StatusBadRequest, invalid.StatusCode)
|
|
require.Contains(t, invalid.Body, "invalid_grant")
|
|
}
|
|
|
|
func TestRefreshIDCTokenInvalidGrantReturnsTypedError(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/token", r.URL.Path)
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
_, _ = w.Write([]byte(`{"error":"invalid_grant","message":"Invalid refresh token provided"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
previous := oidcEndpointOverride
|
|
oidcEndpointOverride = server.URL
|
|
t.Cleanup(func() { oidcEndpointOverride = previous })
|
|
|
|
_, err := RefreshIDCToken(context.Background(), "", "client-id", "client-secret", "revoked-refresh-token", "us-east-1", BuilderIDStartURL)
|
|
require.Error(t, err)
|
|
|
|
var invalid *RefreshTokenInvalidError
|
|
require.True(t, errors.As(err, &invalid))
|
|
require.Equal(t, http.StatusBadRequest, invalid.StatusCode)
|
|
require.Contains(t, invalid.Body, "invalid_grant")
|
|
}
|
|
|
|
func TestExchangeIDCAuthCodePreservesProfileArn(t *testing.T) {
|
|
const profileArn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/EXCHANGE"
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/token":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"accessToken":"access-token","refreshToken":"refresh-token","profileArn":"` + profileArn + `","expiresIn":3600}`))
|
|
case "/userinfo":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"email":"kiro@example.com"}`))
|
|
default:
|
|
t.Fatalf("unexpected path: %s", r.URL.Path)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
previous := oidcEndpointOverride
|
|
oidcEndpointOverride = server.URL
|
|
t.Cleanup(func() { oidcEndpointOverride = previous })
|
|
|
|
token, err := ExchangeIDCAuthCode(context.Background(), "", "client-id", "client-secret", "code", "verifier", "http://127.0.0.1:9876/oauth/callback", "us-east-1", BuilderIDStartURL)
|
|
require.NoError(t, err)
|
|
require.Equal(t, profileArn, token.ProfileArn)
|
|
require.Equal(t, "kiro@example.com", token.Email)
|
|
}
|
|
|
|
func TestRefreshIDCTokenPreservesProfileArn(t *testing.T) {
|
|
const profileArn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/REFRESH"
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/token":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"accessToken":"access-token","refreshToken":"refresh-token","profileArn":"` + profileArn + `","expiresIn":3600}`))
|
|
case "/userinfo":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"email":"kiro@example.com"}`))
|
|
default:
|
|
t.Fatalf("unexpected path: %s", r.URL.Path)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
previous := oidcEndpointOverride
|
|
oidcEndpointOverride = server.URL
|
|
t.Cleanup(func() { oidcEndpointOverride = previous })
|
|
|
|
token, err := RefreshIDCToken(context.Background(), "", "client-id", "client-secret", "refresh-token", "us-east-1", BuilderIDStartURL)
|
|
require.NoError(t, err)
|
|
require.Equal(t, profileArn, token.ProfileArn)
|
|
require.Equal(t, "kiro@example.com", token.Email)
|
|
}
|