Function TestMCPServerUserTokens

enterprise/dbcrypt/dbcrypt_internal_test.go:1471–1573

(t *testing.T)

}

func TestMCPServerUserTokens(t *testing.T) {
	t.Parallel()
	ctx := context.Background()

	const (
		accessToken  = "access-token-value"
		refreshToken = "refresh-token-value"
	)

	// insertConfigAndToken creates a user, an MCP server config, and a
	// user token through the encrypted store.
	insertConfigAndToken := func(
		t *testing.T,
		crypt *dbCrypt,
		ciphers []Cipher,
	) (database.MCPServerConfig, database.MCPServerUserToken) {
		t.Helper()
		user := dbgen.User(t, crypt, database.User{})
		cfg := dbgen.MCPServerConfig(t, crypt, database.MCPServerConfig{
			DisplayName: "Token Test MCP",
			AuthType:    "oauth2",
			CreatedBy:   uuid.NullUUID{UUID: user.ID, Valid: true},
			UpdatedBy:   uuid.NullUUID{UUID: user.ID, Valid: true},
		})

		tok, err := crypt.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{
			MCPServerConfigID: cfg.ID,
			UserID:            user.ID,
			AccessToken:       accessToken,
			RefreshToken:      refreshToken,
			TokenType:         "Bearer",
		})
		require.NoError(t, err)
		require.Equal(t, accessToken, tok.AccessToken)
		require.Equal(t, refreshToken, tok.RefreshToken)
		require.Equal(t, ciphers[0].HexDigest(), tok.AccessTokenKeyID.String)
		require.Equal(t, ciphers[0].HexDigest(), tok.RefreshTokenKeyID.String)
		return cfg, tok
	}

	t.Run("UpsertMCPServerUserToken", func(t *testing.T) {
		t.Parallel()
		db, crypt, ciphers := setup(t)
		cfg, tok := insertConfigAndToken(t, crypt, ciphers)

		// Verify the raw DB values are encrypted.
		rawTok, err := db.GetMCPServerUserToken(ctx, database.GetMCPServerUserTokenParams{
			MCPServerConfigID: cfg.ID,
			UserID:            tok.UserID,
		})
		require.NoError(t, err)
		requireEncryptedEquals(t, ciphers[0], rawTok.AccessToken, accessToken)
		requireEncryptedEquals(t, ciphers[0], rawTok.RefreshToken, refreshToken)
	})

	t.Run("GetMCPServerUserToken", func(t *testing.T) {
		t.Parallel()
		db, crypt, ciphers := setup(t)