Fix WithExtra

kylecarbs · kylecarbs · commit 4d62ebaad5d7 · 2023-10-09T22:11:49.000Z
diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -68,7 +68,7 @@ type FakeIDP struct {
 	// "Authorized Redirect URLs". This can be used to emulate that.
 	hookValidRedirectURL func(redirectURL string) error
 	hookUserInfo         func(email string) (jwt.MapClaims, error)
-	hookExtra            func(email string) map[string]interface{}
+	hookMutateToken      func(token map[string]interface{})
 	fakeCoderd           func(req *http.Request) (*http.Response, error)
 	hookOnRefresh        func(email string) error
 	// Custom authentication for the client. This is useful if you want
@@ -115,9 +115,9 @@ func WithRefresh(hook func(email string) error) func(*FakeIDP) {
 
 // WithExtra returns extra fields that be accessed on the returned Oauth Token.
 // These extra fields can override the default fields (id_token, access_token, etc).
-func WithExtra(extra func(email string) map[string]interface{}) func(*FakeIDP) {
+func WithMutateToken(mutateToken func(token map[string]interface{})) func(*FakeIDP) {
 	return func(f *FakeIDP) {
-		f.hookExtra = extra
+		f.hookMutateToken = mutateToken
 	}
 }
 
@@ -630,10 +630,8 @@ func (f *FakeIDP) httpHandler(t testing.TB) http.Handler {
 			"expires_in":    int64((time.Minute * 5).Seconds()),
 			"id_token":      f.encodeClaims(t, claims),
 		}
-		if f.hookExtra != nil {
-			for k, v := range f.hookExtra(email) {
-				token[k] = v
-			}
+		if f.hookMutateToken != nil {
+			f.hookMutateToken(token)
 		}
 		// Store the claims for the next refresh
 		f.refreshIDTokenClaims.Store(refreshToken, claims)
diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go
@@ -267,11 +267,9 @@ func TestRefreshToken(t *testing.T) {
 		db := dbfake.New()
 		fake, config, link := setupOauth2Test(t, testConfig{
 			FakeIDPOpts: []oidctest.FakeIDPOpt{
-				oidctest.WithExtra(func(email string) map[string]interface{} {
-					return map[string]interface{}{
-						"authed_user": map[string]interface{}{
-							"access_token": "slack-user-token",
-						},
+				oidctest.WithMutateToken(func(token map[string]interface{}) {
+					token["authed_user"] = map[string]interface{}{
+						"access_token": "slack-user-token",
 					}
 				}),
 			},
