package configstore import ( "bytes" "context" "encoding/json" "fmt" "log" "os" "strconv" "strings" "testing" "time" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore/tables" "github.com/maximhq/bifrost/framework/encrypt" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" ) // postgresDSN matches the postgres service in tests/docker-compose.yml and // framework/docker-compose.yml. const postgresDSN = "host=localhost user=bifrost password=bifrost_password dbname=bifrost port=5432 sslmode=disable" // namedDB pairs a backend name with its GORM connection for use in subtests. type namedDB struct { name string db *gorm.DB } // setupTestDB creates an in-memory SQLite database for testing func setupTestDB(t *testing.T) *gorm.DB { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err, "Failed to create test database") // Create the MCP clients table err = db.AutoMigrate(&tables.TableMCPClient{}) require.NoError(t, err, "Failed to migrate test database") return db } // captureLogOutput captures log output during a function execution func captureLogOutput(fn func()) string { var buf bytes.Buffer log.SetOutput(&buf) defer log.SetOutput(os.Stderr) fn() return buf.String() } func TestNormalizeName(t *testing.T) { tests := []struct { name string input string expected string }{ { name: "hyphen to underscore", input: "my-tool", expected: "my_tool", }, { name: "space to underscore", input: "my tool", expected: "my_tool", }, { name: "multiple hyphens", input: "my-super-tool", expected: "my_super_tool", }, { name: "multiple spaces", input: "my super tool", expected: "my_super_tool", }, { name: "leading digits removed", input: "123tool", expected: "tool", }, { name: "leading digits with hyphen", input: "123my-tool", expected: "my_tool", }, { name: "empty after normalization", input: "123", expected: "mcp_client", }, { name: "no change needed", input: "my_tool", expected: "my_tool", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { normalized := normalizeMCPClientName(tt.input) assert.Equal(t, tt.expected, normalized, "normalizeMCPClientName should produce expected output") }) } } func TestFindUniqueName_NoCollision(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // Create a test client with a unique name client := &tables.TableMCPClient{ Name: "existing_client", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err := db.WithContext(ctx).Create(client).Error require.NoError(t, err) // Test findUniqueName with a different base name (no collision) logOutput := captureLogOutput(func() { uniqueName, err := findUniqueNameForTest("new_client", "new_client", 999, db.WithContext(ctx)) require.NoError(t, err) assert.Equal(t, "new_client", uniqueName, "Should return base name when no collision") }) // Should not log anything when there's no collision assert.Empty(t, logOutput, "Should not log when name is available without suffix") } func TestFindUniqueName_WithCollision(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // Create existing clients that will cause collisions // First client with base name client1 := &tables.TableMCPClient{ Name: "my_tool", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err := db.WithContext(ctx).Create(client1).Error require.NoError(t, err) // Second client with first suffix client2 := &tables.TableMCPClient{ Name: "my_tool1", ClientID: "client-2", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err = db.WithContext(ctx).Create(client2).Error require.NoError(t, err) // Test findUniqueName with collision - should find "my_tool2" // excludeID is set to a non-existent ID (999) so all existing clients are considered var uniqueName string logOutput := captureLogOutput(func() { uniqueName, err = findUniqueNameForTest("my_tool", "my-tool", 999, db.WithContext(ctx)) }) require.NoError(t, err) assert.Equal(t, "my_tool2", uniqueName, "Should return name with suffix when collision occurs") assert.Contains(t, logOutput, "MCP Client Name Normalized: 'my-tool' -> 'my_tool2'", "Should log the transformation") } func TestFindUniqueName_MultipleCollisions(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // Create existing clients that will cause multiple collisions client1 := &tables.TableMCPClient{ Name: "test_tool", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err := db.WithContext(ctx).Create(client1).Error require.NoError(t, err) client2 := &tables.TableMCPClient{ Name: "test_tool1", ClientID: "client-2", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err = db.WithContext(ctx).Create(client2).Error require.NoError(t, err) client3 := &tables.TableMCPClient{ Name: "test_tool2", ClientID: "client-3", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err = db.WithContext(ctx).Create(client3).Error require.NoError(t, err) // Test findUniqueName with multiple collisions - should find "test_tool3" var uniqueName string logOutput := captureLogOutput(func() { uniqueName, err = findUniqueNameForTest("test_tool", "test tool", 999, db.WithContext(ctx)) }) require.NoError(t, err) assert.Equal(t, "test_tool3", uniqueName, "Should return name with correct suffix after multiple collisions") assert.Contains(t, logOutput, "MCP Client Name Normalized: 'test tool' -> 'test_tool3'", "Should log the transformation") } func TestFindUniqueName_NormalizationAndCollision(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // Create existing client with normalized name client := &tables.TableMCPClient{ Name: "my_tool", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), } err := db.WithContext(ctx).Create(client).Error require.NoError(t, err) // Test that "my-tool" normalizes to "my_tool" and then collides, requiring suffix var uniqueName string logOutput := captureLogOutput(func() { uniqueName, err = findUniqueNameForTest("my_tool", "my-tool", 999, db.WithContext(ctx)) }) require.NoError(t, err) assert.Equal(t, "my_tool2", uniqueName, "Should handle normalization and collision") assert.Contains(t, logOutput, "MCP Client Name Normalized: 'my-tool' -> 'my_tool2'", "Should log the full transformation") } func TestFindUniqueName_MultipleNormalizationsToSameBase(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // Test case: 3 entries that normalize to the same base name: // "mcp client" -> "mcp_client" // "mcp-client" -> "mcp_client" (collision, becomes "mcp_client2") // "1mcp-client" -> "mcp_client" (collision, becomes "mcp_client3") // Note: In the actual migration, names are processed sequentially and each checks // against all previously created names. To simulate this, we need to create clients // with the original names first, then normalize them in sequence. // Helper function to normalize (same logic as in migrations.go) normalizeName := func(name string) string { normalized := strings.ReplaceAll(name, "-", "_") normalized = strings.ReplaceAll(normalized, " ", "_") normalized = strings.TrimLeftFunc(normalized, func(r rune) bool { return r >= '0' && r <= '9' }) if normalized == "" { normalized = "mcp_client" } return normalized } // Create three clients with original names (simulating pre-migration state) clients := []*tables.TableMCPClient{ { Name: "mcp client", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, { Name: "mcp-client", ClientID: "client-2", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, { Name: "1mcp-client", ClientID: "client-3", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, } for _, client := range clients { err := db.WithContext(ctx).Create(client).Error require.NoError(t, err) } // Now simulate the migration: process each client sequentially // First: "mcp client" -> "mcp_client" (no collision) client1 := clients[0] normalizedName1 := normalizeName(client1.Name) var uniqueName1 string var err error logOutput1 := captureLogOutput(func() { uniqueName1, err = findUniqueNameForTest(normalizedName1, client1.Name, client1.ID, db.WithContext(ctx)) }) require.NoError(t, err) assert.Equal(t, "mcp_client", uniqueName1, "First normalization should use base name") assert.Empty(t, logOutput1, "Should not log when name is available without suffix") // Update first client err = db.WithContext(ctx).Model(client1).Update("name", uniqueName1).Error require.NoError(t, err) // Second: "mcp-client" -> "mcp_client" (collision with "mcp_client", becomes "mcp_client2") // Note: We need to check that "mcp_client" exists (from client1), so it should skip to "mcp_client2" client2 := clients[1] normalizedName2 := normalizeName(client2.Name) var uniqueName2 string logOutput2 := captureLogOutput(func() { uniqueName2, err = findUniqueNameForTest(normalizedName2, client2.Name, client2.ID, db.WithContext(ctx)) }) require.NoError(t, err) // With the updated implementation, suffixes start from 2 when base name exists // So "mcp-client" normalizes to "mcp_client" which collides, becomes "mcp_client2" assert.Equal(t, "mcp_client2", uniqueName2, "Second normalization should get suffix 2 (skipping 1)") assert.Contains(t, logOutput2, "MCP Client Name Normalized: 'mcp-client' -> 'mcp_client2'", "Should log the transformation") // Update second client err = db.WithContext(ctx).Model(client2).Update("name", uniqueName2).Error require.NoError(t, err) // Third: "1mcp-client" -> "mcp_client" (collision with "mcp_client" and "mcp_client2", becomes "mcp_client3") client3 := clients[2] normalizedName3 := normalizeName(client3.Name) var uniqueName3 string logOutput3 := captureLogOutput(func() { uniqueName3, err = findUniqueNameForTest(normalizedName3, client3.Name, client3.ID, db.WithContext(ctx)) }) require.NoError(t, err) // Third normalization finds "mcp_client" and "mcp_client2" exist, so becomes "mcp_client3" assert.Equal(t, "mcp_client3", uniqueName3, "Third normalization should get suffix 3") assert.Contains(t, logOutput3, "MCP Client Name Normalized: '1mcp-client' -> 'mcp_client3'", "Should log the transformation") // Update third client err = db.WithContext(ctx).Model(client3).Update("name", uniqueName3).Error require.NoError(t, err) // Final verification: all three should exist with correct names var finalClients []tables.TableMCPClient err = db.WithContext(ctx).Find(&finalClients).Error require.NoError(t, err) assert.Len(t, finalClients, 3, "Should have all 3 clients") names := make([]string, len(finalClients)) for i, c := range finalClients { names[i] = c.Name } assert.Contains(t, names, "mcp_client", "Should contain mcp_client") assert.Contains(t, names, "mcp_client2", "Should contain mcp_client2") assert.Contains(t, names, "mcp_client3", "Should contain mcp_client3") } func TestFindUniqueName_MigrationScenarioWithInMemoryTracking(t *testing.T) { db := setupTestDB(t) ctx := context.Background() // This test simulates the exact migration scenario where clients are processed in a loop // and we need to track assigned names in memory to avoid transaction visibility issues // Create three clients with original names (simulating pre-migration state) clients := []*tables.TableMCPClient{ { Name: "mcp client", ClientID: "client-1", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, { Name: "mcp-client", ClientID: "client-2", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, { Name: "1mcp-client", ClientID: "client-3", ConnectionType: "stdio", CreatedAt: time.Now(), UpdatedAt: time.Now(), }, } for _, client := range clients { err := db.WithContext(ctx).Create(client).Error require.NoError(t, err) } // Simulate the migration: process clients in a loop with in-memory tracking assignedNames := make(map[string]bool) normalizeName := func(name string) string { normalized := strings.ReplaceAll(name, "-", "_") normalized = strings.ReplaceAll(normalized, " ", "_") normalized = strings.TrimLeftFunc(normalized, func(r rune) bool { return r >= '0' && r <= '9' }) if normalized == "" { normalized = "mcp_client" } return normalized } var logOutputs []string for _, client := range clients { originalName := client.Name needsUpdate := strings.Contains(originalName, "-") || strings.Contains(originalName, " ") || (len(originalName) > 0 && originalName[0] >= '0' && originalName[0] <= '9') if needsUpdate { normalizedName := normalizeName(originalName) uniqueName, err := findUniqueNameForTestWithTracking(normalizedName, originalName, client.ID, db.WithContext(ctx), assignedNames) require.NoError(t, err) // Capture log output logOutput := captureLogOutput(func() { // Log if name changed if originalName != uniqueName { log.Printf("MCP Client Name Normalized: '%s' -> '%s'", originalName, uniqueName) } }) if logOutput != "" { logOutputs = append(logOutputs, logOutput) } // Update client err = db.WithContext(ctx).Model(client).Update("name", uniqueName).Error require.NoError(t, err) } } // Verify all three clients have correct names var finalClients []tables.TableMCPClient err := db.WithContext(ctx).Find(&finalClients).Error require.NoError(t, err) assert.Len(t, finalClients, 3, "Should have all 3 clients") names := make([]string, len(finalClients)) for i, c := range finalClients { names[i] = c.Name } assert.Contains(t, names, "mcp_client", "Should contain mcp_client") assert.Contains(t, names, "mcp_client2", "Should contain mcp_client2") assert.Contains(t, names, "mcp_client3", "Should contain mcp_client3") // Verify logging: should log all three transformations allLogs := strings.Join(logOutputs, "") assert.Contains(t, allLogs, "MCP Client Name Normalized: 'mcp client' -> 'mcp_client'", "Should log first normalization") assert.Contains(t, allLogs, "MCP Client Name Normalized: 'mcp-client' -> 'mcp_client2'", "Should log second normalization") assert.Contains(t, allLogs, "MCP Client Name Normalized: '1mcp-client' -> 'mcp_client3'", "Should log third normalization") } // findUniqueNameForTestWithTracking is a test helper that tracks assigned names in memory func findUniqueNameForTestWithTracking(baseName string, originalName string, excludeID uint, tx *gorm.DB, assignedNames map[string]bool) (string, error) { // First check if base name is already assigned in this migration if !assignedNames[baseName] { // Also check database for existing names (excluding current client) var count int64 err := tx.Model(&tables.TableMCPClient{}).Where("name = ? AND id != ?", baseName, excludeID).Count(&count).Error if err != nil { return "", fmt.Errorf("failed to check name availability: %w", err) } if count == 0 { // Name is available assignedNames[baseName] = true // Log normalization even when no collision if originalName != baseName { log.Printf("MCP Client Name Normalized: '%s' -> '%s'", originalName, baseName) } return baseName, nil } } // Name exists (either assigned in this migration or in database), try with number suffix starting from 2 suffix := 2 const maxSuffix = 1000 for { if suffix > maxSuffix { return "", fmt.Errorf("could not find unique name after %d attempts for base name: %s", maxSuffix, baseName) } candidateName := baseName + strconv.Itoa(suffix) // Check both in-memory map and database if !assignedNames[candidateName] { var count int64 err := tx.Model(&tables.TableMCPClient{}).Where("name = ? AND id != ?", candidateName, excludeID).Count(&count).Error if err != nil { return "", fmt.Errorf("failed to check name availability: %w", err) } if count == 0 { // Found available name assignedNames[candidateName] = true log.Printf("MCP Client Name Normalized: '%s' -> '%s'", originalName, candidateName) return candidateName, nil } } suffix++ } } // findUniqueNameForTest is a test helper that extracts the findUniqueName logic // This mirrors the implementation in migrations.go for testing func findUniqueNameForTest(baseName string, originalName string, excludeID uint, tx *gorm.DB) (string, error) { // First, try the base name var count int64 err := tx.Model(&tables.TableMCPClient{}).Where("name = ? AND id != ?", baseName, excludeID).Count(&count).Error if err != nil { return "", fmt.Errorf("failed to check name availability: %w", err) } if count == 0 { // Name is available return baseName, nil } // Name exists, try with number suffix starting from 2 // (base name is conceptually "1", so collisions start from "2") suffix := 2 const maxSuffix = 1000 // Safety limit to prevent infinite loops for { if suffix > maxSuffix { return "", fmt.Errorf("could not find unique name after %d attempts for base name: %s", maxSuffix, baseName) } candidateName := baseName + strconv.Itoa(suffix) err := tx.Model(&tables.TableMCPClient{}).Where("name = ? AND id != ?", candidateName, excludeID).Count(&count).Error if err != nil { return "", fmt.Errorf("failed to check name availability: %w", err) } if count == 0 { // Found available name - log the transformation log.Printf("MCP Client Name Normalized: '%s' -> '%s'", originalName, candidateName) return candidateName, nil } suffix++ } } // setupProviderTestDBWithoutStoreRawColumn creates an in-memory SQLite database // with the config_providers table but WITHOUT the store_raw_request_response column, // simulating the pre-migration state. func setupProviderTestDBWithoutStoreRawColumn(t *testing.T) *gorm.DB { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err, "Failed to create test database") // Create the config_providers table manually without store_raw_request_response column // This simulates the pre-migration state err = db.Exec(` CREATE TABLE config_providers ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(50) NOT NULL UNIQUE, network_config_json TEXT, concurrency_buffer_json TEXT, proxy_config_json TEXT, custom_provider_config_json TEXT, pricing_overrides_json TEXT, send_back_raw_request BOOLEAN DEFAULT 0, send_back_raw_response BOOLEAN DEFAULT 0, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, budget_id VARCHAR(255), rate_limit_id VARCHAR(255), config_hash VARCHAR(255), status VARCHAR(50) DEFAULT 'unknown', description TEXT, encryption_status VARCHAR(20) DEFAULT 'plain_text' ) `).Error require.NoError(t, err, "Failed to create config_providers table") // Create the migrations table for the migrator (matches migrator.DefaultOptions.TableName) err = db.Exec(` CREATE TABLE IF NOT EXISTS migrations ( id VARCHAR(255) PRIMARY KEY ) `).Error require.NoError(t, err, "Failed to create migrations table") return db } // trySetupPostgresDBWithoutStoreRawColumn attempts to connect to Postgres and creates // the config_providers table WITHOUT the store_raw_request_response column. // Returns nil (without skipping the test) if Postgres is unavailable. func trySetupPostgresDBWithoutStoreRawColumn(t *testing.T, testSuffix string) *gorm.DB { t.Helper() db, err := gorm.Open(postgres.Open(postgresDSN), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) if err != nil { return nil } // Verify the connection is actually live before proceeding. sqlDB, err := db.DB() if err != nil { return nil } if err := sqlDB.Ping(); err != nil { return nil } // Drop config_providers to start fresh (for this specific test). // Use CASCADE to drop dependent objects (composite types, sequences, etc.). db.Exec("DROP TABLE IF EXISTS config_providers CASCADE") // Clear migration tracking without dropping the table — other test packages // (e.g. logstore) may share this Postgres instance and use the same table // concurrently. CREATE IF NOT EXISTS is safe even if the table already // exists from a previous test or a concurrent package. db.Exec(`CREATE TABLE IF NOT EXISTS migrations ( id VARCHAR(255) PRIMARY KEY )`) db.Exec("DELETE FROM migrations") // Create the config_providers table manually without store_raw_request_response column // This simulates the pre-migration state (PostgreSQL syntax) err = db.Exec(` CREATE TABLE config_providers ( id SERIAL PRIMARY KEY, name VARCHAR(50) NOT NULL UNIQUE, network_config_json TEXT, concurrency_buffer_json TEXT, proxy_config_json TEXT, custom_provider_config_json TEXT, pricing_overrides_json TEXT, send_back_raw_request BOOLEAN DEFAULT false, send_back_raw_response BOOLEAN DEFAULT false, created_at TIMESTAMP NOT NULL, updated_at TIMESTAMP NOT NULL, budget_id VARCHAR(255), rate_limit_id VARCHAR(255), config_hash VARCHAR(255), status VARCHAR(50) DEFAULT 'unknown', description TEXT, encryption_status VARCHAR(20) DEFAULT 'plain_text' ) `).Error if err != nil { return nil } // Clean up after the test — drop config_providers but leave migrations // intact for concurrent test packages. t.Cleanup(func() { db.Exec("DELETE FROM migrations") db.Exec("DROP TABLE IF EXISTS config_providers CASCADE") }) return db } // forEachProviderMigrationDB returns backends for provider migration tests. // Always includes SQLite; includes Postgres when available. func forEachProviderMigrationDB(t *testing.T, testSuffix string) []namedDB { t.Helper() dbs := []namedDB{{"sqlite", setupProviderTestDBWithoutStoreRawColumn(t)}} if pgDB := trySetupPostgresDBWithoutStoreRawColumn(t, testSuffix); pgDB != nil { dbs = append(dbs, namedDB{"postgres", pgDB}) } return dbs } func TestMigrationAddStoreRawRequestResponseColumn(t *testing.T) { tests := []struct { name string sendBackRawRequest bool sendBackRawResponse bool expectedStoreRawRequestResponse bool }{ { name: "both false - store should be false", sendBackRawRequest: false, sendBackRawResponse: false, expectedStoreRawRequestResponse: false, }, { name: "request true response false - store should be true", sendBackRawRequest: true, sendBackRawResponse: false, expectedStoreRawRequestResponse: true, }, { name: "request false response true - store should be true", sendBackRawRequest: false, sendBackRawResponse: true, expectedStoreRawRequestResponse: true, }, { name: "both true - store should be true", sendBackRawRequest: true, sendBackRawResponse: true, expectedStoreRawRequestResponse: true, }, } for _, tt := range tests { tt := tt t.Run(tt.name, func(t *testing.T) { for _, ndb := range forEachProviderMigrationDB(t, tt.name) { ndb := ndb t.Run(ndb.name, func(t *testing.T) { db := ndb.db ctx := context.Background() providerName := "test_provider" staleHash := "stale_hash_before_migration" now := time.Now() // Insert a provider with the old schema (no store_raw_request_response column) err := db.Exec(` INSERT INTO config_providers ( name, send_back_raw_request, send_back_raw_response, config_hash, created_at, updated_at, encryption_status ) VALUES (?, ?, ?, ?, ?, ?, ?) `, providerName, tt.sendBackRawRequest, tt.sendBackRawResponse, staleHash, now, now, "plain_text").Error require.NoError(t, err, "Failed to insert test provider") // Verify column does not exist before migration hasColumn := db.Migrator().HasColumn(&tables.TableProvider{}, "store_raw_request_response") assert.False(t, hasColumn, "store_raw_request_response column should not exist before migration") // Run the migration err = migrationAddStoreRawRequestResponseColumn(ctx, db) require.NoError(t, err, "Migration should succeed") // Verify column exists after migration hasColumn = db.Migrator().HasColumn(&tables.TableProvider{}, "store_raw_request_response") assert.True(t, hasColumn, "store_raw_request_response column should exist after migration") // Fetch the provider and verify values var result struct { Name string SendBackRawRequest bool SendBackRawResponse bool StoreRawRequestResponse bool ConfigHash string } err = db.Table("config_providers"). Select("name, send_back_raw_request, send_back_raw_response, store_raw_request_response, config_hash"). Where("name = ?", providerName). Scan(&result).Error require.NoError(t, err, "Failed to fetch provider after migration") // Verify store_raw_request_response was set correctly assert.Equal(t, tt.expectedStoreRawRequestResponse, result.StoreRawRequestResponse, "store_raw_request_response should be set based on send_back_raw_request OR send_back_raw_response") // Verify config_hash was updated (not the stale hash) assert.NotEqual(t, staleHash, result.ConfigHash, "config_hash should be updated after migration") // Verify the hash matches what GenerateConfigHash would produce expectedConfig := ProviderConfig{ SendBackRawRequest: tt.sendBackRawRequest, SendBackRawResponse: tt.sendBackRawResponse, StoreRawRequestResponse: tt.expectedStoreRawRequestResponse, } expectedHash, err := expectedConfig.GenerateConfigHash(providerName) require.NoError(t, err, "Failed to generate expected hash") assert.Equal(t, expectedHash, result.ConfigHash, "config_hash should match the expected hash from GenerateConfigHash") }) } }) } } func TestMigrationAddStoreRawRequestResponseColumn_MultipleProviders(t *testing.T) { for _, ndb := range forEachProviderMigrationDB(t, "multiple") { ndb := ndb t.Run(ndb.name, func(t *testing.T) { db := ndb.db ctx := context.Background() now := time.Now() // Insert multiple providers with different configurations providers := []struct { name string sendBackRawRequest bool sendBackRawResponse bool }{ {"provider_neither", false, false}, {"provider_request_only", true, false}, {"provider_response_only", false, true}, {"provider_both", true, true}, } for _, p := range providers { err := db.Exec(` INSERT INTO config_providers ( name, send_back_raw_request, send_back_raw_response, config_hash, created_at, updated_at, encryption_status ) VALUES (?, ?, ?, ?, ?, ?, ?) `, p.name, p.sendBackRawRequest, p.sendBackRawResponse, "stale_hash", now, now, "plain_text").Error require.NoError(t, err, "Failed to insert provider %s", p.name) } // Run the migration err := migrationAddStoreRawRequestResponseColumn(ctx, db) require.NoError(t, err, "Migration should succeed") // Verify each provider for _, p := range providers { var result struct { StoreRawRequestResponse bool ConfigHash string } err := db.Table("config_providers"). Select("store_raw_request_response, config_hash"). Where("name = ?", p.name). Scan(&result).Error require.NoError(t, err, "Failed to fetch provider %s", p.name) expectedStore := p.sendBackRawRequest || p.sendBackRawResponse assert.Equal(t, expectedStore, result.StoreRawRequestResponse, "Provider %s: store_raw_request_response mismatch", p.name) assert.NotEqual(t, "stale_hash", result.ConfigHash, "Provider %s: config_hash should be updated", p.name) } }) } } func TestMigrationAddStoreRawRequestResponseColumn_Idempotent(t *testing.T) { for _, ndb := range forEachProviderMigrationDB(t, "idempotent") { ndb := ndb t.Run(ndb.name, func(t *testing.T) { db := ndb.db ctx := context.Background() now := time.Now() providerName := "idempotent_test_provider" // Insert a provider err := db.Exec(` INSERT INTO config_providers ( name, send_back_raw_request, send_back_raw_response, config_hash, created_at, updated_at, encryption_status ) VALUES (?, ?, ?, ?, ?, ?, ?) `, providerName, true, false, "stale_hash", now, now, "plain_text").Error require.NoError(t, err, "Failed to insert test provider") // Run the migration first time err = migrationAddStoreRawRequestResponseColumn(ctx, db) require.NoError(t, err, "First migration should succeed") // Get the hash after first migration var firstHash string err = db.Table("config_providers"). Select("config_hash"). Where("name = ?", providerName). Scan(&firstHash).Error require.NoError(t, err) // Run the migration second time (should be idempotent) err = migrationAddStoreRawRequestResponseColumn(ctx, db) require.NoError(t, err, "Second migration should succeed (idempotent)") // Verify hash is unchanged after second run var secondHash string err = db.Table("config_providers"). Select("config_hash"). Where("name = ?", providerName). Scan(&secondHash).Error require.NoError(t, err) assert.Equal(t, firstHash, secondHash, "Hash should remain unchanged after idempotent migration run") }) } } // setupKeyDBWithLegacyDeploymentColumns creates an in-memory SQLite database // with the config_keys table including legacy deployment columns and NO aliases_json. // This simulates the pre-aliases, post-encryption database state. func setupKeyDBWithLegacyDeploymentColumns(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) // Create migrations table (required by the migrator) err = db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error require.NoError(t, err) // Create config_keys table with legacy deployment columns and encryption_status // but WITHOUT aliases_json — simulating a DB from before the aliases migration err = db.Exec(` CREATE TABLE config_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(255) NOT NULL UNIQUE, provider_id INTEGER NOT NULL, provider VARCHAR(50), key_id VARCHAR(255) NOT NULL UNIQUE, value TEXT NOT NULL, models_json TEXT, blacklisted_models_json TEXT, weight REAL, enabled BOOLEAN DEFAULT true, created_at TIMESTAMP NOT NULL, updated_at TIMESTAMP NOT NULL, config_hash VARCHAR(255), azure_endpoint TEXT, azure_api_version TEXT, azure_client_id TEXT, azure_client_secret TEXT, azure_tenant_id TEXT, azure_scopes TEXT, azure_deployments_json TEXT, vertex_project_id TEXT, vertex_project_number TEXT, vertex_region TEXT, vertex_auth_credentials TEXT, vertex_deployments_json TEXT, bedrock_access_key TEXT, bedrock_secret_key TEXT, bedrock_session_token TEXT, bedrock_region TEXT, bedrock_arn TEXT, bedrock_role_arn TEXT, bedrock_external_id TEXT, bedrock_role_session_name TEXT, bedrock_batch_s3_config TEXT, bedrock_deployments_json TEXT, replicate_deployments_json TEXT, replicate_use_deployments_endpoint BOOLEAN, vllm_url TEXT, vllm_model_name TEXT, use_for_batch_api BOOLEAN DEFAULT false, status VARCHAR(50) DEFAULT 'unknown', description TEXT, encryption_status VARCHAR(20) DEFAULT 'plain_text', ollama_url TEXT, sgl_url TEXT ) `).Error require.NoError(t, err) t.Cleanup(func() { db.Exec("DELETE FROM migrations") db.Exec("DROP TABLE IF EXISTS config_keys") }) return db } // TestMigrationDropDeploymentColumnsAndAddAliases_EncryptedRows tests that the // aliases migration correctly handles rows where encryption_status is 'encrypted' // but deployment columns contain plaintext JSON (because those columns were never // in the encryption list). This was the root cause of the base64 decode crash. func TestMigrationDropDeploymentColumnsAndAddAliases_EncryptedRows(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled, skipping encrypted rows test") } db := setupKeyDBWithLegacyDeploymentColumns(t) ctx := context.Background() now := time.Now().UTC().Format("2006-01-02 15:04:05") // Encrypt the key values (simulating what EncryptPlaintextRows would have done) encVal1, err := encrypt.Encrypt("sk-azure-key-1") require.NoError(t, err) encVal2, err := encrypt.Encrypt("sk-vertex-key-1") require.NoError(t, err) encVal3, err := encrypt.Encrypt("sk-replicate-key-1") require.NoError(t, err) // Insert encrypted rows with PLAINTEXT deployment JSON. // azure_deployments_json, vertex_deployments_json, and replicate_deployments_json // were never in the encryption list, so they remain plaintext even though the row // is marked as encrypted. err = db.Exec(` INSERT INTO config_keys ( name, provider_id, provider, key_id, value, models_json, azure_deployments_json, encryption_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, 'encrypted', ?, ?)`, "azure-key-1", 1, "azure", "ak-1", encVal1, `["*"]`, `{"gpt-4":"dep-gpt4","gpt-3.5":"dep-gpt35"}`, now, now, ).Error require.NoError(t, err) err = db.Exec(` INSERT INTO config_keys ( name, provider_id, provider, key_id, value, models_json, vertex_deployments_json, encryption_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, 'encrypted', ?, ?)`, "vertex-key-1", 2, "vertex", "vk-1", encVal2, `["*"]`, `{"gemini-pro":"dep-gemini"}`, now, now, ).Error require.NoError(t, err) err = db.Exec(` INSERT INTO config_keys ( name, provider_id, provider, key_id, value, models_json, replicate_deployments_json, encryption_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, 'encrypted', ?, ?)`, "replicate-key-1", 3, "replicate", "rk-1", encVal3, `["*"]`, `{"llama":"dep-llama"}`, now, now, ).Error require.NoError(t, err) // Run the aliases migration — this should NOT crash with base64 decode error. err = migrationDropDeploymentColumnsAndAddAliases(ctx, db) require.NoError(t, err, "migration should not crash on encrypted rows with plaintext deployment data") // Verify aliases_json was populated and is readable via GORM hooks var keys []tables.TableKey err = db.Order("name").Find(&keys).Error require.NoError(t, err) require.Len(t, keys, 3) // Azure key should have its deployment data migrated to aliases assert.Equal(t, "azure-key-1", keys[0].Name) assert.NotNil(t, keys[0].Aliases) // Vertex key assert.Equal(t, "vertex-key-1", keys[2].Name) assert.NotNil(t, keys[2].Aliases) // Replicate key — sorted alphabetically, it's between azure and vertex assert.Equal(t, "replicate-key-1", keys[1].Name) assert.NotNil(t, keys[1].Aliases) } // TestMigrationDropDeploymentColumnsAndAddAliases_BedrockEncrypted tests the branch // where bedrock_deployments_json is already encrypted before migration. The migration // should detect it's already encrypted (Decrypt succeeds) and NOT double-encrypt it. func TestMigrationDropDeploymentColumnsAndAddAliases_BedrockEncrypted(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled, skipping bedrock encrypted test") } db := setupKeyDBWithLegacyDeploymentColumns(t) ctx := context.Background() now := time.Now().UTC().Format("2006-01-02 15:04:05") // Encrypt the key value encVal, err := encrypt.Encrypt("sk-bedrock-key-1") require.NoError(t, err) // Encrypt the deployments JSON (simulating bedrock which WAS in the encryption list) bedrockDeployments := `{"claude":"dep-claude","claude-instant":"dep-instant"}` encDeployments, err := encrypt.Encrypt(bedrockDeployments) require.NoError(t, err) // Insert a bedrock row where bedrock_deployments_json is already encrypted err = db.Exec(` INSERT INTO config_keys ( name, provider_id, provider, key_id, value, models_json, bedrock_deployments_json, encryption_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, 'encrypted', ?, ?)`, "bedrock-key-1", 4, "bedrock", "bk-1", encVal, `["*"]`, encDeployments, now, now, ).Error require.NoError(t, err) // Run the aliases migration — should detect already-encrypted data and skip re-encryption err = migrationDropDeploymentColumnsAndAddAliases(ctx, db) require.NoError(t, err, "migration should handle already-encrypted bedrock deployments") // Verify aliases_json was populated and is readable via GORM hooks (AfterFind decrypts) var keys []tables.TableKey err = db.Find(&keys).Error require.NoError(t, err) require.Len(t, keys, 1) assert.Equal(t, "bedrock-key-1", keys[0].Name) assert.NotNil(t, keys[0].Aliases) // Verify the aliases contain the original deployment data (not double-encrypted) aliases := keys[0].Aliases assert.Contains(t, aliases, "claude") assert.Equal(t, "dep-claude", aliases["claude"]) assert.Equal(t, "dep-instant", aliases["claude-instant"]) } // ============================================================================ // Helper: full migration DB setup // ============================================================================ // setupFullMigrationDB creates a fresh in-memory SQLite database and runs the // full triggerMigrations chain (the same code path as production startup). // Returns an RDBConfigStore for CRUD verification and the raw *gorm.DB for // low-level assertions. // testDBCounter ensures each test gets a unique shared-cache in-memory SQLite URI. var testDBCounter int64 func setupFullMigrationDB(t *testing.T) (*RDBConfigStore, *gorm.DB) { t.Helper() // Use a unique shared-cache URI so all connections see the same in-memory DB // without requiring MaxOpenConns(1) (which deadlocks when code opens transactions // and queries on s.db concurrently). n := time.Now().UnixNano() + testDBCounter testDBCounter++ dsn := fmt.Sprintf("file:testdb_%d?mode=memory&cache=shared", n) db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err, "Failed to create test database") ctx := context.Background() err = triggerMigrations(ctx, db) require.NoError(t, err, "triggerMigrations should succeed on a fresh DB") store := &RDBConfigStore{logger: bifrost.NewDefaultLogger(schemas.LogLevelInfo)} store.db.Store(db) store.migrateOnFreshFn = func(ctx context.Context, fn func(context.Context, *gorm.DB) error) error { return fn(ctx, store.DB()) } store.refreshPoolFn = func(ctx context.Context) error { return nil } return store, db } // ============================================================================ // Part 2: Full-chain integration tests // ============================================================================ func TestTriggerMigrations_FreshDB(t *testing.T) { _, db := setupFullMigrationDB(t) // Every critical table should exist after the full migration chain. criticalTables := []interface{}{ &tables.TableProvider{}, &tables.TableKey{}, &tables.TableVirtualKey{}, &tables.TableMCPClient{}, &tables.TableBudget{}, &tables.TableRateLimit{}, &tables.TablePlugin{}, &tables.TableCustomer{}, &tables.TableTeam{}, &tables.SessionsTable{}, &tables.TableOauthConfig{}, &tables.TableOauthToken{}, &tables.TableModelPricing{}, &tables.TableGovernanceConfig{}, &tables.TableClientConfig{}, &tables.TableVirtualKeyProviderConfig{}, &tables.TableVirtualKeyMCPConfig{}, } migrator := db.Migrator() for _, table := range criticalTables { assert.True(t, migrator.HasTable(table), "table should exist: %T", table) } } func TestTriggerMigrations_Idempotent(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) ctx := context.Background() // First run err = triggerMigrations(ctx, db) require.NoError(t, err, "first triggerMigrations should succeed") // Second run – must be a no-op, not an error. err = triggerMigrations(ctx, db) require.NoError(t, err, "second triggerMigrations should succeed (idempotent)") // Tables should still be intact assert.True(t, db.Migrator().HasTable(&tables.TableProvider{}), "TableProvider should still exist") assert.True(t, db.Migrator().HasTable(&tables.TableKey{}), "TableKey should still exist") assert.True(t, db.Migrator().HasTable(&tables.TableVirtualKey{}), "TableVirtualKey should still exist") } func TestFullMigration_ProviderAndKeyCRUD(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled") } store, db := setupFullMigrationDB(t) ctx := context.Background() config := ProviderConfig{ Keys: []schemas.Key{ { ID: "key-uuid-1", Name: "openai-primary", Value: *schemas.NewEnvVar("sk-test-secret-key-12345"), Models: schemas.WhiteList{"*"}, Weight: 1.0, }, }, } err := store.AddProvider(ctx, "openai", config) require.NoError(t, err) // Read back through store (AfterFind hooks decrypt) result, err := store.GetProvidersConfig(ctx) require.NoError(t, err) require.Len(t, result, 1) require.Contains(t, result, schemas.ModelProvider("openai")) openai := result[schemas.ModelProvider("openai")] require.Len(t, openai.Keys, 1) assert.Equal(t, "openai-primary", openai.Keys[0].Name) assert.Equal(t, "sk-test-secret-key-12345", openai.Keys[0].Value.GetValue()) // Verify encryption at the raw DB level var rawValue string var rawStatus string err = db.Table("config_keys").Select("value").Where("key_id = ?", "key-uuid-1").Scan(&rawValue).Error require.NoError(t, err) err = db.Table("config_keys").Select("encryption_status").Where("key_id = ?", "key-uuid-1").Scan(&rawStatus).Error require.NoError(t, err) assert.NotEqual(t, "sk-test-secret-key-12345", rawValue, "key value should be encrypted at rest") assert.Equal(t, "encrypted", rawStatus) } func TestFullMigration_VirtualKeyCRUD(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled") } store, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now() vk := &tables.TableVirtualKey{ ID: "vk-test-001", Name: "test-virtual-key", Value: "vk-secret-value-12345", IsActive: true, CreatedAt: now, UpdatedAt: now, } err := store.CreateVirtualKey(ctx, vk) require.NoError(t, err) // Read back vks, err := store.GetVirtualKeys(ctx) require.NoError(t, err) require.Len(t, vks, 1) assert.Equal(t, "vk-test-001", vks[0].ID) assert.Equal(t, "test-virtual-key", vks[0].Name) assert.Equal(t, "vk-secret-value-12345", vks[0].Value) // AfterFind decrypts assert.True(t, vks[0].IsActive) // Verify encryption at raw DB level var rawValue, rawStatus, rawHash string err = db.Table("governance_virtual_keys"). Select("value").Where("id = ?", "vk-test-001").Scan(&rawValue).Error require.NoError(t, err) err = db.Table("governance_virtual_keys"). Select("encryption_status").Where("id = ?", "vk-test-001").Scan(&rawStatus).Error require.NoError(t, err) err = db.Table("governance_virtual_keys"). Select("value_hash").Where("id = ?", "vk-test-001").Scan(&rawHash).Error require.NoError(t, err) assert.NotEqual(t, "vk-secret-value-12345", rawValue, "VK value should be encrypted at rest") assert.Equal(t, "encrypted", rawStatus) assert.NotEmpty(t, rawHash, "value_hash should be populated") } func TestFullMigration_MCPClientCRUD(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled") } store, db := setupFullMigrationDB(t) ctx := context.Background() clientConfig := &schemas.MCPClientConfig{ ID: "mcp-client-001", Name: "test_mcp_server", ConnectionType: schemas.MCPConnectionTypeSSE, ConnectionString: schemas.NewEnvVar("https://mcp.example.com/sse"), ToolsToExecute: schemas.WhiteList{"*"}, } err := store.CreateMCPClientConfig(ctx, clientConfig) require.NoError(t, err) // Read back through store mcpClient, err := store.GetMCPClientByName(ctx, "test_mcp_server") require.NoError(t, err) require.NotNil(t, mcpClient) assert.Equal(t, "mcp-client-001", mcpClient.ClientID) assert.Equal(t, "test_mcp_server", mcpClient.Name) assert.Equal(t, "sse", mcpClient.ConnectionType) assert.Equal(t, "https://mcp.example.com/sse", mcpClient.ConnectionString.GetValue()) // Verify encryption at raw DB level var rawConnStr, rawStatus string err = db.Table("config_mcp_clients"). Select("connection_string").Where("client_id = ?", "mcp-client-001").Scan(&rawConnStr).Error require.NoError(t, err) err = db.Table("config_mcp_clients"). Select("encryption_status").Where("client_id = ?", "mcp-client-001").Scan(&rawStatus).Error require.NoError(t, err) assert.NotEqual(t, "https://mcp.example.com/sse", rawConnStr, "connection string should be encrypted at rest") assert.Equal(t, "encrypted", rawStatus) } func TestFullMigration_EncryptPlaintextRows(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled") } store, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now().UTC().Format("2006-01-02 15:04:05") // Insert a provider first (FK for config_keys) err := db.Exec(`INSERT INTO config_providers (name, encryption_status, created_at, updated_at) VALUES (?, 'plain_text', ?, ?)`, "openai", now, now).Error require.NoError(t, err) // Get the provider ID var providerID uint err = db.Table("config_providers").Select("id").Where("name = ?", "openai").Scan(&providerID).Error require.NoError(t, err) // Insert plaintext key (bypassing GORM hooks) err = db.Exec(`INSERT INTO config_keys (name, provider_id, provider, key_id, value, models_json, encryption_status, created_at, updated_at) VALUES (?, ?, 'openai', ?, ?, '["*"]', 'plain_text', ?, ?)`, "plaintext-key", providerID, "pk-1", "sk-plaintext-secret", now, now).Error require.NoError(t, err) // Insert plaintext virtual key err = db.Exec(`INSERT INTO governance_virtual_keys (id, name, value, is_active, encryption_status, created_at, updated_at) VALUES (?, ?, ?, true, 'plain_text', ?, ?)`, "vk-plain-1", "plaintext-vk", "vk-plain-secret", now, now).Error require.NoError(t, err) // Verify they are plaintext before encryption var preStatus string err = db.Table("config_keys").Select("encryption_status").Where("key_id = ?", "pk-1").Scan(&preStatus).Error require.NoError(t, err) assert.Equal(t, "plain_text", preStatus) // Run the encryption pass err = store.EncryptPlaintextRows(ctx) require.NoError(t, err) // Verify raw DB: encryption_status changed, values differ var rawKeyValue, rawKeyStatus string err = db.Table("config_keys").Select("value").Where("key_id = ?", "pk-1").Scan(&rawKeyValue).Error require.NoError(t, err) err = db.Table("config_keys").Select("encryption_status").Where("key_id = ?", "pk-1").Scan(&rawKeyStatus).Error require.NoError(t, err) assert.NotEqual(t, "sk-plaintext-secret", rawKeyValue, "key should be encrypted after EncryptPlaintextRows") assert.Equal(t, "encrypted", rawKeyStatus) var rawVKValue, rawVKStatus string err = db.Table("governance_virtual_keys").Select("value").Where("id = ?", "vk-plain-1").Scan(&rawVKValue).Error require.NoError(t, err) err = db.Table("governance_virtual_keys").Select("encryption_status").Where("id = ?", "vk-plain-1").Scan(&rawVKStatus).Error require.NoError(t, err) assert.NotEqual(t, "vk-plain-secret", rawVKValue, "VK should be encrypted after EncryptPlaintextRows") assert.Equal(t, "encrypted", rawVKStatus) // Verify GORM read decrypts correctly var key tables.TableKey err = db.Where("key_id = ?", "pk-1").First(&key).Error require.NoError(t, err) assert.Equal(t, "sk-plaintext-secret", key.Value.GetValue()) var vk tables.TableVirtualKey err = db.Where("id = ?", "vk-plain-1").First(&vk).Error require.NoError(t, err) assert.Equal(t, "vk-plain-secret", vk.Value) } func TestFullMigration_EndToEnd(t *testing.T) { if !encrypt.IsEnabled() { t.Skip("encryption not enabled") } store, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now() // Add two providers with keys for _, p := range []struct { provider string keyID string keyName string keyValue string }{ {"openai", "key-oa-1", "openai-key", "sk-openai-secret"}, {"anthropic", "key-ant-1", "anthropic-key", "sk-anthropic-secret"}, } { err := store.AddProvider(ctx, schemas.ModelProvider(p.provider), ProviderConfig{ Keys: []schemas.Key{{ ID: p.keyID, Name: p.keyName, Value: *schemas.NewEnvVar(p.keyValue), Models: schemas.WhiteList{"*"}, Weight: 1.0, }}, }) require.NoError(t, err, "AddProvider %s", p.provider) } // Add virtual keys for _, vk := range []struct { id, name, value string }{ {"vk-1", "vk-alpha", "vk-alpha-secret"}, {"vk-2", "vk-beta", "vk-beta-secret"}, } { err := store.CreateVirtualKey(ctx, &tables.TableVirtualKey{ ID: vk.id, Name: vk.name, Value: vk.value, IsActive: true, CreatedAt: now, UpdatedAt: now, }) require.NoError(t, err, "CreateVirtualKey %s", vk.name) } // Add MCP client err := store.CreateMCPClientConfig(ctx, &schemas.MCPClientConfig{ ID: "mcp-e2e-1", Name: "e2e_mcp_client", ConnectionType: schemas.MCPConnectionTypeSSE, ConnectionString: schemas.NewEnvVar("https://mcp.e2e.test/sse"), ToolsToExecute: schemas.WhiteList{"*"}, }) require.NoError(t, err) // Verify providers providers, err := store.GetProvidersConfig(ctx) require.NoError(t, err) assert.Len(t, providers, 2) assert.Contains(t, providers, schemas.ModelProvider("openai")) assert.Contains(t, providers, schemas.ModelProvider("anthropic")) assert.Equal(t, "sk-openai-secret", providers["openai"].Keys[0].Value.GetValue()) assert.Equal(t, "sk-anthropic-secret", providers["anthropic"].Keys[0].Value.GetValue()) // Verify virtual keys vks, err := store.GetVirtualKeys(ctx) require.NoError(t, err) assert.Len(t, vks, 2) // Verify MCP client mcpClient, err := store.GetMCPClientByName(ctx, "e2e_mcp_client") require.NoError(t, err) assert.Equal(t, "https://mcp.e2e.test/sse", mcpClient.ConnectionString.GetValue()) // Verify all sensitive data is encrypted at the raw DB level type encCheck struct { table, column, whereCol, whereVal, plaintext string } checks := []encCheck{ {"config_keys", "value", "key_id", "key-oa-1", "sk-openai-secret"}, {"config_keys", "value", "key_id", "key-ant-1", "sk-anthropic-secret"}, {"governance_virtual_keys", "value", "id", "vk-1", "vk-alpha-secret"}, {"governance_virtual_keys", "value", "id", "vk-2", "vk-beta-secret"}, {"config_mcp_clients", "connection_string", "client_id", "mcp-e2e-1", "https://mcp.e2e.test/sse"}, } for _, c := range checks { var rawVal string err := db.Table(c.table).Select(c.column).Where(fmt.Sprintf("%s = ?", c.whereCol), c.whereVal).Scan(&rawVal).Error require.NoError(t, err) assert.NotEqual(t, c.plaintext, rawVal, "raw %s.%s for %s=%s should be encrypted", c.table, c.column, c.whereCol, c.whereVal) } } // ============================================================================ // Part 3: Individual complex migration tests // ============================================================================ // setupPreEncryptionDB runs all migrations up to (but not including) the // encryption columns migration. This is approximated by running migrationInit // plus the essential early migrations on a fresh SQLite DB. func setupPreEncryptionDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) ctx := context.Background() // Run the initial migration to create core tables err = migrationInit(ctx, db) require.NoError(t, err) // Sessions table is created by a later migration, but required by the // encryption columns migration. err = migrationAddSessionsTable(ctx, db) require.NoError(t, err) return db } func TestMigrationAddEncryptionColumns(t *testing.T) { db := setupPreEncryptionDB(t) ctx := context.Background() now := time.Now().UTC().Format("2006-01-02 15:04:05") // Insert rows BEFORE encryption columns exist (they won't have encryption_status) err := db.Exec(`INSERT INTO config_providers (name, created_at, updated_at) VALUES (?, ?, ?)`, "openai", now, now).Error require.NoError(t, err) var providerID uint err = db.Table("config_providers").Select("id").Where("name = ?", "openai").Scan(&providerID).Error require.NoError(t, err) err = db.Exec(`INSERT INTO config_keys (name, provider_id, provider, key_id, value, created_at, updated_at) VALUES (?, ?, 'openai', ?, ?, ?, ?)`, "test-key", providerID, "ek-1", "sk-test", now, now).Error require.NoError(t, err) err = db.Exec(`INSERT INTO governance_virtual_keys (id, name, value, is_active, created_at, updated_at) VALUES (?, ?, ?, true, ?, ?)`, "vk-enc-1", "enc-vk", "vk-value", now, now).Error require.NoError(t, err) err = db.Exec(`INSERT INTO sessions (token, created_at, updated_at, expires_at) VALUES (?, ?, ?, ?)`, "sess-token", now, now, now).Error require.NoError(t, err) // Run the encryption columns migration err = migrationAddEncryptionColumns(ctx, db) require.NoError(t, err) // Verify encryption_status column exists and is backfilled on all 9 tables encTables := []string{ "config_keys", "governance_virtual_keys", "sessions", "oauth_configs", "oauth_tokens", "config_mcp_clients", "config_providers", "config_vector_store", "config_plugins", } for _, table := range encTables { var count int64 err := db.Table(table).Where("encryption_status = 'plain_text'").Count(&count).Error if err != nil { // Table might be empty, just verify column exists assert.True(t, db.Migrator().HasColumn(table, "encryption_status"), "encryption_status column should exist on %s", table) continue } } // Verify pre-existing rows have encryption_status = 'plain_text' var keyStatus string err = db.Table("config_keys").Select("encryption_status").Where("key_id = ?", "ek-1").Scan(&keyStatus).Error require.NoError(t, err) assert.Equal(t, "plain_text", keyStatus) var providerStatus string err = db.Table("config_providers").Select("encryption_status").Where("name = ?", "openai").Scan(&providerStatus).Error require.NoError(t, err) assert.Equal(t, "plain_text", providerStatus) var vkStatus string err = db.Table("governance_virtual_keys").Select("encryption_status").Where("id = ?", "vk-enc-1").Scan(&vkStatus).Error require.NoError(t, err) assert.Equal(t, "plain_text", vkStatus) // Verify value_hash on governance_virtual_keys is NULL (not empty string) var rawHash *string err = db.Table("governance_virtual_keys").Select("value_hash").Where("id = ?", "vk-enc-1").Scan(&rawHash).Error require.NoError(t, err) assert.Nil(t, rawHash, "value_hash should be NULL, not empty string") // Verify token_hash on sessions is NULL var tokenHash *string err = db.Table("sessions").Select("token_hash").Where("token = ?", "sess-token").Scan(&tokenHash).Error require.NoError(t, err) assert.Nil(t, tokenHash, "token_hash should be NULL, not empty string") // Idempotency: running again should not error err = migrationAddEncryptionColumns(ctx, db) require.NoError(t, err) } func TestMigrationCleanupMCPClientToolsConfig(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) ctx := context.Background() // Create the MCP clients table using AutoMigrate err = db.AutoMigrate(&tables.TableMCPClient{}) require.NoError(t, err) // Create the migrations tracking table err = db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error require.NoError(t, err) now := time.Now() // Insert clients with various tools_to_execute_json states clients := []struct { name string tools string }{ {"client_empty_array", "[]"}, {"client_empty_string", ""}, {"client_wildcard", `["*"]`}, {"client_specific", `["tool1","tool2"]`}, } for i, c := range clients { err := db.Exec(`INSERT INTO config_mcp_clients (client_id, name, connection_type, tools_to_execute_json, created_at, updated_at, encryption_status) VALUES (?, ?, 'stdio', ?, ?, ?, 'plain_text')`, fmt.Sprintf("client-%d", i), c.name, c.tools, now, now).Error require.NoError(t, err) } // Also insert one with NULL tools err = db.Exec(`INSERT INTO config_mcp_clients (client_id, name, connection_type, tools_to_execute_json, created_at, updated_at, encryption_status) VALUES (?, ?, 'stdio', NULL, ?, ?, 'plain_text')`, "client-null", "client_null_tools", now, now).Error require.NoError(t, err) // Run the cleanup migration err = migrationCleanupMCPClientToolsConfig(ctx, db) require.NoError(t, err) // Verify: empty/null → ["*"], existing values preserved for _, tc := range []struct { name string expected string }{ {"client_empty_array", `["*"]`}, {"client_empty_string", `["*"]`}, {"client_null_tools", `["*"]`}, {"client_wildcard", `["*"]`}, {"client_specific", `["tool1","tool2"]`}, } { var tools string err := db.Table("config_mcp_clients").Select("tools_to_execute_json"). Where("name = ?", tc.name).Scan(&tools).Error require.NoError(t, err) assert.Equal(t, tc.expected, tools, "tools for %s", tc.name) } } func TestMigrationAddConfigHashColumn(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) ctx := context.Background() // Create migrations table err = db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error require.NoError(t, err) // Create tables WITHOUT config_hash column to simulate pre-migration state err = db.Exec(`CREATE TABLE config_providers ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(50) NOT NULL UNIQUE, network_config_json TEXT, concurrency_buffer_json TEXT, proxy_config_json TEXT, custom_provider_config_json TEXT, send_back_raw_request BOOLEAN DEFAULT 0, send_back_raw_response BOOLEAN DEFAULT 0, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, encryption_status VARCHAR(20) DEFAULT 'plain_text' )`).Error require.NoError(t, err) err = db.Exec(`CREATE TABLE config_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(255) NOT NULL, key_id VARCHAR(255) NOT NULL UNIQUE, provider_id INTEGER NOT NULL, provider VARCHAR(50), value TEXT NOT NULL, models_json TEXT, weight REAL, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, encryption_status VARCHAR(20) DEFAULT 'plain_text' )`).Error require.NoError(t, err) now := time.Now() // Insert provider and key without config_hash err = db.Exec(`INSERT INTO config_providers (name, created_at, updated_at, encryption_status) VALUES (?, ?, ?, 'plain_text')`, "openai", now, now).Error require.NoError(t, err) var providerID uint err = db.Table("config_providers").Select("id").Where("name = ?", "openai").Scan(&providerID).Error require.NoError(t, err) err = db.Exec(`INSERT INTO config_keys (name, key_id, provider_id, provider, value, models_json, created_at, updated_at, encryption_status) VALUES (?, ?, ?, 'openai', ?, '["*"]', ?, ?, 'plain_text')`, "test-key", "ck-1", providerID, "sk-test-value", now, now).Error require.NoError(t, err) // Run the migration err = migrationAddConfigHashColumn(ctx, db) require.NoError(t, err) // Verify config_hash column exists assert.True(t, db.Migrator().HasColumn(&tables.TableProvider{}, "config_hash")) assert.True(t, db.Migrator().HasColumn(&tables.TableKey{}, "config_hash")) // Verify hashes are non-empty var providerHash string err = db.Table("config_providers").Select("config_hash").Where("name = ?", "openai").Scan(&providerHash).Error require.NoError(t, err) assert.NotEmpty(t, providerHash, "provider config_hash should be backfilled") var keyHash string err = db.Table("config_keys").Select("config_hash").Where("key_id = ?", "ck-1").Scan(&keyHash).Error require.NoError(t, err) assert.NotEmpty(t, keyHash, "key config_hash should be backfilled") // Verify provider hash matches expected computation expectedProviderConfig := ProviderConfig{} expectedHash, err := expectedProviderConfig.GenerateConfigHash("openai") require.NoError(t, err) assert.Equal(t, expectedHash, providerHash, "provider hash should match GenerateConfigHash output") } func TestMigrationBackfillEmptyVirtualKeyConfigs(t *testing.T) { // Use full migration DB to get all tables, then insert test data store, db := setupFullMigrationDB(t) _ = store ctx := context.Background() now := time.Now() // Clear the migration tracking for the specific migration so it runs again db.Exec(`DELETE FROM migrations WHERE id = 'backfill_empty_virtual_key_configs'`) // Create a provider err := db.Exec(`INSERT INTO config_providers (name, encryption_status, created_at, updated_at) VALUES ('openai', 'plain_text', ?, ?)`, now, now).Error require.NoError(t, err) // Create a virtual key with NO provider configs err = db.Exec(`INSERT INTO governance_virtual_keys (id, name, value, is_active, encryption_status, created_at, updated_at) VALUES ('vk-backfill-1', 'backfill-vk', 'vk-value', true, 'plain_text', ?, ?)`, now, now).Error require.NoError(t, err) // Verify no provider configs exist var provConfigCount int64 db.Table("governance_virtual_key_provider_configs").Where("virtual_key_id = ?", "vk-backfill-1").Count(&provConfigCount) assert.Equal(t, int64(0), provConfigCount, "should have no provider configs before migration") // Run the migration err = migrationBackfillEmptyVirtualKeyConfigs(ctx, db) require.NoError(t, err) // Verify provider configs were created for the VK db.Table("governance_virtual_key_provider_configs").Where("virtual_key_id = ?", "vk-backfill-1").Count(&provConfigCount) assert.Equal(t, int64(1), provConfigCount, "should have 1 provider config after backfill (one for openai)") // Verify the provider config has correct defaults var allowAllKeys bool err = db.Table("governance_virtual_key_provider_configs"). Select("allow_all_keys").Where("virtual_key_id = ?", "vk-backfill-1").Scan(&allowAllKeys).Error require.NoError(t, err) assert.True(t, allowAllKeys, "backfilled provider config should have AllowAllKeys=true") // Verify config_hash was recomputed (non-empty) var vkHash string err = db.Table("governance_virtual_keys").Select("config_hash"). Where("id = ?", "vk-backfill-1").Scan(&vkHash).Error require.NoError(t, err) assert.NotEmpty(t, vkHash, "VK config_hash should be recomputed after backfill") } func TestMigrationBackfillAllowedModelsWildcard(t *testing.T) { _, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now() // Clear migration tracking so it runs again db.Exec(`DELETE FROM migrations WHERE id = 'backfill_allowed_models_wildcard'`) // Create a provider err := db.Exec(`INSERT INTO config_providers (name, encryption_status, created_at, updated_at) VALUES ('openai', 'plain_text', ?, ?)`, now, now).Error require.NoError(t, err) var providerID uint err = db.Table("config_providers").Select("id").Where("name = ?", "openai").Scan(&providerID).Error require.NoError(t, err) // Create a key with empty models_json err = db.Exec(`INSERT INTO config_keys (name, key_id, provider_id, provider, value, models_json, encryption_status, created_at, updated_at) VALUES ('empty-models-key', 'emk-1', ?, 'openai', 'sk-test', '[]', 'plain_text', ?, ?)`, providerID, now, now).Error require.NoError(t, err) // Create a VK err = db.Exec(`INSERT INTO governance_virtual_keys (id, name, value, is_active, encryption_status, created_at, updated_at) VALUES ('vk-wildcard-1', 'wildcard-vk', 'vk-val', true, 'plain_text', ?, ?)`, now, now).Error require.NoError(t, err) // Create a provider config with empty allowed_models err = db.Exec(`INSERT INTO governance_virtual_key_provider_configs (virtual_key_id, provider, allowed_models, allow_all_keys) VALUES ('vk-wildcard-1', 'openai', '[]', true)`).Error require.NoError(t, err) // Run the migration err = migrationBackfillAllowedModelsWildcard(ctx, db) require.NoError(t, err) // Verify provider config allowed_models changed to ["*"] var allowedModels string err = db.Table("governance_virtual_key_provider_configs"). Select("allowed_models").Where("virtual_key_id = ?", "vk-wildcard-1").Scan(&allowedModels).Error require.NoError(t, err) assert.Equal(t, `["*"]`, allowedModels, "allowed_models should be backfilled to wildcard") // Verify key models_json changed to ["*"] var modelsJSON string err = db.Table("config_keys").Select("models_json").Where("key_id = ?", "emk-1").Scan(&modelsJSON).Error require.NoError(t, err) assert.Equal(t, `["*"]`, modelsJSON, "models_json should be backfilled to wildcard") // Verify key config_hash was recomputed var keyHash string err = db.Table("config_keys").Select("config_hash").Where("key_id = ?", "emk-1").Scan(&keyHash).Error require.NoError(t, err) assert.NotEmpty(t, keyHash, "key config_hash should be recomputed") } func TestMigrationRemoveServerPrefixFromMCPTools(t *testing.T) { _, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now() // Clear migration tracking so it runs again db.Exec(`DELETE FROM migrations WHERE id = 'remove_server_prefix_from_mcp_tools'`) // Insert an MCP client with prefixed tool names toolsJSON, _ := json.Marshal([]string{"my_server_tool1", "my_server_tool2", "standalone_tool"}) autoToolsJSON, _ := json.Marshal([]string{"my_server_auto1"}) err := db.Exec(`INSERT INTO config_mcp_clients (client_id, name, connection_type, tools_to_execute_json, tools_to_auto_execute_json, encryption_status, created_at, updated_at) VALUES (?, ?, 'stdio', ?, ?, 'plain_text', ?, ?)`, "mcp-prefix-1", "my_server", string(toolsJSON), string(autoToolsJSON), now, now).Error require.NoError(t, err) // Run the migration err = migrationRemoveServerPrefixFromMCPTools(ctx, db) require.NoError(t, err) // Verify tools had prefixes stripped var resultToolsJSON string err = db.Table("config_mcp_clients").Select("tools_to_execute_json"). Where("client_id = ?", "mcp-prefix-1").Scan(&resultToolsJSON).Error require.NoError(t, err) var resultTools []string err = json.Unmarshal([]byte(resultToolsJSON), &resultTools) require.NoError(t, err) assert.Contains(t, resultTools, "tool1", "should have stripped my_server_ prefix") assert.Contains(t, resultTools, "tool2", "should have stripped my_server_ prefix") assert.Contains(t, resultTools, "standalone_tool", "should preserve non-prefixed tool") assert.NotContains(t, resultTools, "my_server_tool1", "original prefixed name should be gone") // Verify auto-execute tools var resultAutoToolsJSON string err = db.Table("config_mcp_clients").Select("tools_to_auto_execute_json"). Where("client_id = ?", "mcp-prefix-1").Scan(&resultAutoToolsJSON).Error require.NoError(t, err) var resultAutoTools []string err = json.Unmarshal([]byte(resultAutoToolsJSON), &resultAutoTools) require.NoError(t, err) assert.Contains(t, resultAutoTools, "auto1", "should have stripped my_server_ prefix from auto tools") } func TestMigrationRemoveServerPrefixFromMCPTools_Collision(t *testing.T) { _, db := setupFullMigrationDB(t) ctx := context.Background() now := time.Now() // Clear migration tracking db.Exec(`DELETE FROM migrations WHERE id = 'remove_server_prefix_from_mcp_tools'`) // Client where stripping the prefix would cause a collision: // "srv_read" (prefixed) → "read", but "read" already exists in the list toolsJSON, _ := json.Marshal([]string{"srv_read", "read"}) err := db.Exec(`INSERT INTO config_mcp_clients (client_id, name, connection_type, tools_to_execute_json, encryption_status, created_at, updated_at) VALUES (?, ?, 'stdio', ?, 'plain_text', ?, ?)`, "mcp-collision", "srv", string(toolsJSON), now, now).Error require.NoError(t, err) // Run the migration — should not error, collision is handled err = migrationRemoveServerPrefixFromMCPTools(ctx, db) require.NoError(t, err) // The collision drops the duplicate, keeping "read" once var resultJSON string err = db.Table("config_mcp_clients").Select("tools_to_execute_json"). Where("client_id = ?", "mcp-collision").Scan(&resultJSON).Error require.NoError(t, err) var resultTools []string err = json.Unmarshal([]byte(resultJSON), &resultTools) require.NoError(t, err) assert.Contains(t, resultTools, "read") // Should not have duplicates readCount := 0 for _, tool := range resultTools { if tool == "read" { readCount++ } } assert.Equal(t, 1, readCount, "should deduplicate on collision") } func TestMigrationReplaceEnableLiteLLMWithCompatColumns(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err) ctx := context.Background() // Create migrations table err = db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error require.NoError(t, err) // Create config_client from the GORM model so DropColumn works (SQLite recreates // the table from model columns). Then add the legacy column on top. err = db.AutoMigrate(&tables.TableClientConfig{}) require.NoError(t, err) // Add the legacy column that the migration will read and drop err = db.Exec(`ALTER TABLE config_client ADD COLUMN enable_litellm_fallbacks BOOLEAN DEFAULT 0`).Error require.NoError(t, err) now := time.Now() // Insert a row with enable_litellm_fallbacks = true err = db.Exec(`INSERT INTO config_client (enable_litellm_fallbacks, created_at, updated_at) VALUES (1, ?, ?)`, now, now).Error require.NoError(t, err) // Insert a row with enable_litellm_fallbacks = false err = db.Exec(`INSERT INTO config_client (enable_litellm_fallbacks, created_at, updated_at) VALUES (0, ?, ?)`, now, now).Error require.NoError(t, err) // Run the migration err = migrationReplaceEnableLiteLLMWithCompatColumns(ctx, db) require.NoError(t, err) // Verify new columns exist mgr := db.Migrator() assert.True(t, mgr.HasColumn(&tables.TableClientConfig{}, "compat_convert_text_to_chat")) assert.True(t, mgr.HasColumn(&tables.TableClientConfig{}, "compat_convert_chat_to_responses")) assert.True(t, mgr.HasColumn(&tables.TableClientConfig{}, "compat_should_drop_params")) assert.True(t, mgr.HasColumn(&tables.TableClientConfig{}, "compat_should_convert_params")) // Verify data migration: row 1 had litellm=true → compat_convert_text_to_chat=true type compatRow struct { ID uint CompatConvertTextToChat bool `gorm:"column:compat_convert_text_to_chat"` CompatShouldConvertParams bool `gorm:"column:compat_should_convert_params"` } var rows []compatRow err = db.Table("config_client"). Select("id, compat_convert_text_to_chat, compat_should_convert_params"). Order("id").Find(&rows).Error require.NoError(t, err) require.Len(t, rows, 2) assert.True(t, rows[0].CompatConvertTextToChat, "row with litellm=true should have compat_convert_text_to_chat=true") assert.False(t, rows[0].CompatShouldConvertParams, "compat_should_convert_params should default to false") assert.False(t, rows[1].CompatConvertTextToChat, "row with litellm=false should have compat_convert_text_to_chat=false") assert.False(t, rows[1].CompatShouldConvertParams, "compat_should_convert_params should default to false") } // setupCalendarAlignedPreMigrationDB creates a SQLite DB with governance_virtual_keys, // governance_budgets, and governance_rate_limits tables, then drops the calendar_aligned // column from budgets and rate_limits to simulate the pre-migration schema state. func setupCalendarAlignedPreMigrationDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err, "Failed to create test database") require.NoError(t, db.AutoMigrate( &tables.TableBudget{}, &tables.TableRateLimit{}, &tables.TableVirtualKey{}, ), "Failed to auto-migrate governance tables") // Simulate pre-migration state: drop calendar_aligned from budgets and rate_limits. mig := db.Migrator() if mig.HasColumn(&tables.TableBudget{}, "calendar_aligned") { require.NoError(t, mig.DropColumn(&tables.TableBudget{}, "calendar_aligned")) } if mig.HasColumn(&tables.TableRateLimit{}, "calendar_aligned") { require.NoError(t, mig.DropColumn(&tables.TableRateLimit{}, "calendar_aligned")) } require.NoError(t, db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error) return db } // insertBudgetRaw inserts a budget via raw SQL so the pre-migration table (without // calendar_aligned) accepts the row without tripping BeforeSave hooks. func insertBudgetRaw(t *testing.T, db *gorm.DB, id string, vkID *string) { t.Helper() now := time.Now() err := db.Exec(` INSERT INTO governance_budgets (id, max_limit, reset_duration, last_reset, current_usage, virtual_key_id, config_hash, created_at, updated_at) VALUES (?, ?, ?, ?, 0, ?, '', ?, ?) `, id, 100.0, "1d", now, vkID, now, now).Error require.NoError(t, err, "Failed to insert budget %s", id) } // insertRateLimitRaw inserts a rate_limit via raw SQL mirroring insertBudgetRaw. func insertRateLimitRaw(t *testing.T, db *gorm.DB, id string) { t.Helper() now := time.Now() tokenMax := int64(1000) tokenDur := "1h" err := db.Exec(` INSERT INTO governance_rate_limits (id, token_max_limit, token_reset_duration, token_current_usage, token_last_reset, request_current_usage, request_last_reset, config_hash, created_at, updated_at) VALUES (?, ?, ?, 0, ?, 0, ?, '', ?, ?) `, id, tokenMax, tokenDur, now, now, now, now).Error require.NoError(t, err, "Failed to insert rate limit %s", id) } // insertVKRaw inserts a virtual key via raw SQL so we can set rate_limit_id without // running the VK BeforeSave hook (which handles encryption). func insertVKRaw(t *testing.T, db *gorm.DB, id, name, value string, rateLimitID *string, calendarAligned bool) { t.Helper() now := time.Now() err := db.Exec(` INSERT INTO governance_virtual_keys (id, name, value, is_active, rate_limit_id, calendar_aligned, encryption_status, value_hash, config_hash, created_at, updated_at) VALUES (?, ?, ?, 1, ?, ?, 'plain_text', ?, '', ?, ?) `, id, name, value, rateLimitID, calendarAligned, id+"-hash", now, now).Error require.NoError(t, err, "Failed to insert virtual key %s", id) } func setupLegacyBudgetOwnerMigrationDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) require.NoError(t, err, "Failed to create test database") require.NoError(t, db.AutoMigrate( &tables.TableBudget{}, &tables.TableVirtualKey{}, &tables.TableVirtualKeyProviderConfig{}, &tables.TableTeam{}, ), "Failed to auto-migrate legacy budget owner test tables") // Add legacy budget_id columns WITH foreign key constraints, matching production // schema. The FK definitions are what caused SQLite to reject ALTER TABLE DROP // COLUMN and what triggered FK reference corruption via rename propagation. require.NoError(t, db.Exec(`ALTER TABLE governance_virtual_keys ADD COLUMN budget_id VARCHAR(255) REFERENCES governance_budgets(id)`).Error) require.NoError(t, db.Exec(`ALTER TABLE governance_virtual_key_provider_configs ADD COLUMN budget_id VARCHAR(255) REFERENCES governance_budgets(id)`).Error) require.NoError(t, db.Exec(`ALTER TABLE governance_teams ADD COLUMN budget_id VARCHAR(255) REFERENCES governance_budgets(id)`).Error) require.NoError(t, db.Exec(`CREATE TABLE IF NOT EXISTS migrations (id VARCHAR(255) PRIMARY KEY)`).Error) return db } func insertProviderConfigRaw(t *testing.T, db *gorm.DB, id uint, virtualKeyID, provider string) { t.Helper() err := db.Exec(` INSERT INTO governance_virtual_key_provider_configs (id, virtual_key_id, provider, allowed_models, allow_all_keys) VALUES (?, ?, ?, '[]', 1) `, id, virtualKeyID, provider).Error require.NoError(t, err, "Failed to insert provider config %d", id) } func insertTeamRaw(t *testing.T, db *gorm.DB, id, name string) { t.Helper() now := time.Now() err := db.Exec(` INSERT INTO governance_teams (id, name, config_hash, created_at, updated_at) VALUES (?, ?, '', ?, ?) `, id, name, now, now).Error require.NoError(t, err, "Failed to insert team %s", id) } // TestMigrationCalendarAligned_AddColumnsAndBackfill exercises the full migration: // column addition on governance_budgets and governance_rate_limits, plus backfill // of calendar_aligned=true for rows attached to any virtual key. Rows NOT attached // to a VK must remain at the default (false). func TestMigrationCalendarAligned_AddColumnsAndBackfill(t *testing.T) { db := setupCalendarAlignedPreMigrationDB(t) ctx := context.Background() // Pre-condition: columns should be absent before running the migration. mig := db.Migrator() assert.False(t, mig.HasColumn(&tables.TableBudget{}, "calendar_aligned"), "calendar_aligned should NOT exist on budgets before migration") assert.False(t, mig.HasColumn(&tables.TableRateLimit{}, "calendar_aligned"), "calendar_aligned should NOT exist on rate_limits before migration") // Seed: VK-aligned carries the legacy calendar_aligned=true flag; its rate limit // and both attached budgets SHOULD be backfilled. insertRateLimitRaw(t, db, "rl-aligned") vkAlignedID := "vk-aligned" rlAlignedID := "rl-aligned" insertVKRaw(t, db, vkAlignedID, "vk-aligned", "vk-aligned-value", &rlAlignedID, true) insertBudgetRaw(t, db, "budget-aligned-1", &vkAlignedID) insertBudgetRaw(t, db, "budget-aligned-2", &vkAlignedID) // Seed: VK-unaligned carries calendar_aligned=false; its rate limit and // budgets MUST NOT be backfilled (preserves legacy per-VK semantics). insertRateLimitRaw(t, db, "rl-unaligned") vkUnalignedID := "vk-unaligned" rlUnalignedID := "rl-unaligned" insertVKRaw(t, db, vkUnalignedID, "vk-unaligned", "vk-unaligned-value", &rlUnalignedID, false) insertBudgetRaw(t, db, "budget-unaligned", &vkUnalignedID) // Seed: VK-aligned-no-rl has calendar_aligned=true but no rate limit — its // single budget should still be backfilled. vkAlignedNoRLID := "vk-aligned-no-rl" insertVKRaw(t, db, vkAlignedNoRLID, "vk-aligned-no-rl", "vk-aligned-no-rl-value", nil, true) insertBudgetRaw(t, db, "budget-aligned-no-rl", &vkAlignedNoRLID) // Seed: orphaned rows (no VK owner) — these must NOT be backfilled. insertRateLimitRaw(t, db, "rl-orphan") insertBudgetRaw(t, db, "budget-orphan", nil) // Run the migration. require.NoError(t, migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx, db), "migration should succeed") // Post-condition: columns exist. assert.True(t, mig.HasColumn(&tables.TableBudget{}, "calendar_aligned"), "calendar_aligned should exist on budgets after migration") assert.True(t, mig.HasColumn(&tables.TableRateLimit{}, "calendar_aligned"), "calendar_aligned should exist on rate_limits after migration") // Verify per-row state. type budgetRow struct { ID string CalendarAligned bool } type rlRow struct { ID string CalendarAligned bool } assertBudget := func(id string, expected bool) { t.Helper() var row budgetRow err := db.Table("governance_budgets"). Select("id, calendar_aligned"). Where("id = ?", id).Scan(&row).Error require.NoError(t, err) assert.Equal(t, expected, row.CalendarAligned, "budget %s calendar_aligned mismatch", id) } assertRateLimit := func(id string, expected bool) { t.Helper() var row rlRow err := db.Table("governance_rate_limits"). Select("id, calendar_aligned"). Where("id = ?", id).Scan(&row).Error require.NoError(t, err) assert.Equal(t, expected, row.CalendarAligned, "rate_limit %s calendar_aligned mismatch", id) } // Rows attached to VKs with calendar_aligned=true should be backfilled. assertBudget("budget-aligned-1", true) assertBudget("budget-aligned-2", true) assertBudget("budget-aligned-no-rl", true) assertRateLimit("rl-aligned", true) // Rows attached to VKs with calendar_aligned=false must remain false — // the migration must preserve the legacy per-VK semantic. assertBudget("budget-unaligned", false) assertRateLimit("rl-unaligned", false) // Orphaned rows should retain the default (false). assertBudget("budget-orphan", false) assertRateLimit("rl-orphan", false) } // TestMigrationCalendarAligned_StaleRateLimitID verifies that a VK pointing at a // deleted rate_limit row is skipped rather than aborting the migration (the FK // is intentionally not DB-enforced on TableVirtualKey). func TestMigrationCalendarAligned_StaleRateLimitID(t *testing.T) { db := setupCalendarAlignedPreMigrationDB(t) ctx := context.Background() // VK-stale carries calendar_aligned=true (so the migration touches it) and // references a rate_limit_id that doesn't exist in the table. stale := "rl-does-not-exist" insertVKRaw(t, db, "vk-stale", "vk-stale", "vk-stale-value", &stale, true) insertBudgetRaw(t, db, "budget-stale-vk", strPtr("vk-stale")) // VK-ok has calendar_aligned=true and a valid rate_limit so we can verify // the loop keeps processing after the stale skip. insertRateLimitRaw(t, db, "rl-ok") vkOKID := "vk-ok" rlOK := "rl-ok" insertVKRaw(t, db, vkOKID, "vk-ok", "vk-ok-value", &rlOK, true) insertBudgetRaw(t, db, "budget-ok", &vkOKID) // Migration must NOT return an error despite the stale reference. require.NoError(t, migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx, db), "stale rate_limit_id should be skipped, not abort the migration") // The VK with the stale ref still has its budgets backfilled. var staleVKBudget bool require.NoError(t, db.Table("governance_budgets"). Select("calendar_aligned"). Where("id = ?", "budget-stale-vk").Scan(&staleVKBudget).Error) assert.True(t, staleVKBudget, "budget attached to the stale-ref VK should still be migrated") // Non-stale VK should have been fully processed. var okRL bool require.NoError(t, db.Table("governance_rate_limits"). Select("calendar_aligned"). Where("id = ?", "rl-ok").Scan(&okRL).Error) assert.True(t, okRL, "valid rate_limit should be migrated even when a prior VK had a stale ref") } // TestMigrationCalendarAligned_Idempotent ensures the migration can run twice // without error — the migrator library must record the migration ID so the // second invocation is a no-op. func TestMigrationCalendarAligned_Idempotent(t *testing.T) { db := setupCalendarAlignedPreMigrationDB(t) ctx := context.Background() insertRateLimitRaw(t, db, "rl-1") vkID := "vk-1" rl := "rl-1" insertVKRaw(t, db, vkID, "vk-1", "vk-1-value", &rl, true) insertBudgetRaw(t, db, "budget-1", &vkID) require.NoError(t, migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx, db), "first run should succeed") require.NoError(t, migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx, db), "second run should be a no-op") } func TestMigrationAddMultiBudgetTables_DropsLegacyBudgetColumnsAndBackfillsOwners(t *testing.T) { db := setupLegacyBudgetOwnerMigrationDB(t) ctx := context.Background() mig := db.Migrator() vkID := "vk-legacy" insertVKRaw(t, db, vkID, "vk-legacy", "vk-legacy-value", nil, false) insertBudgetRaw(t, db, "budget-vk-legacy", nil) require.NoError(t, db.Exec(`UPDATE governance_virtual_keys SET budget_id = ? WHERE id = ?`, "budget-vk-legacy", vkID).Error) insertProviderConfigRaw(t, db, 1001, vkID, "openai") insertBudgetRaw(t, db, "budget-pc-legacy", nil) require.NoError(t, db.Exec(`UPDATE governance_virtual_key_provider_configs SET budget_id = ? WHERE id = ?`, "budget-pc-legacy", 1001).Error) require.True(t, mig.HasColumn("governance_virtual_keys", "budget_id")) require.True(t, mig.HasColumn("governance_virtual_key_provider_configs", "budget_id")) require.NoError(t, migrationAddMultiBudgetTables(ctx, db)) assert.False(t, mig.HasColumn("governance_virtual_keys", "budget_id"), "vk budget_id should be dropped by migration") assert.False(t, mig.HasColumn("governance_virtual_key_provider_configs", "budget_id"), "provider config budget_id should be dropped by migration") var vkBudgetOwner string require.NoError(t, db.Table("governance_budgets").Select("virtual_key_id").Where("id = ?", "budget-vk-legacy").Scan(&vkBudgetOwner).Error) assert.Equal(t, vkID, vkBudgetOwner, "migration should backfill governance_budgets.virtual_key_id") var providerConfigBudgetOwner uint require.NoError(t, db.Table("governance_budgets").Select("provider_config_id").Where("id = ?", "budget-pc-legacy").Scan(&providerConfigBudgetOwner).Error) assert.Equal(t, uint(1001), providerConfigBudgetOwner, "migration should backfill governance_budgets.provider_config_id") // Verify FK references in governance_budgets still point at the correct // table names — not corrupted backup/temp names from SQLite rename propagation. assertNoCorruptedFKReferences(t, db) } func TestMigrationAddTeamBudgetsToBudgetsTable_DropsLegacyBudgetColumnAndBackfillsOwners(t *testing.T) { db := setupLegacyBudgetOwnerMigrationDB(t) ctx := context.Background() mig := db.Migrator() insertTeamRaw(t, db, "team-legacy", "Legacy Team") insertBudgetRaw(t, db, "budget-team-legacy", nil) require.NoError(t, db.Exec(`UPDATE governance_teams SET budget_id = ? WHERE id = ?`, "budget-team-legacy", "team-legacy").Error) require.True(t, mig.HasColumn("governance_teams", "budget_id")) require.NoError(t, migrationAddTeamBudgetsToBudgetsTable(ctx, db)) assert.False(t, mig.HasColumn("governance_teams", "budget_id"), "team budget_id should be dropped by migration") var teamBudgetOwner string require.NoError(t, db.Table("governance_budgets").Select("team_id").Where("id = ?", "budget-team-legacy").Scan(&teamBudgetOwner).Error) assert.Equal(t, "team-legacy", teamBudgetOwner, "migration should backfill governance_budgets.team_id") // Verify FK references in governance_budgets still point at the correct // table names — not corrupted backup/temp names from SQLite rename propagation. assertNoCorruptedFKReferences(t, db) } // migration is part of the startup chain so a fresh DB emerges with the column // present on both governance_budgets and governance_rate_limits. func TestMigrationCalendarAligned_WiredIntoTriggerMigrations(t *testing.T) { _, db := setupFullMigrationDB(t) mig := db.Migrator() assert.True(t, mig.HasColumn(&tables.TableBudget{}, "calendar_aligned"), "triggerMigrations should add calendar_aligned to governance_budgets") assert.True(t, mig.HasColumn(&tables.TableRateLimit{}, "calendar_aligned"), "triggerMigrations should add calendar_aligned to governance_rate_limits") } // assertNoCorruptedFKReferences checks that no table in the database has FK // references pointing at temporary/backup table names left behind by a // SQLite table rebuild. This is the actual failure mode from the production // bug: ALTER TABLE RENAME propagated into FK definitions in other tables, // leaving references like governance_teams__legacy_budget_backup. func assertNoCorruptedFKReferences(t *testing.T, db *gorm.DB) { t.Helper() type masterRow struct { Name string `gorm:"column:name"` SQL string `gorm:"column:sql"` } var rows []masterRow require.NoError(t, db.Raw(`SELECT name, sql FROM sqlite_master WHERE type = 'table' AND sql IS NOT NULL`).Scan(&rows).Error) for _, row := range rows { assert.NotContains(t, row.SQL, "__dump", "table %s has FK referencing a dump table: %s", row.Name, row.SQL) assert.NotContains(t, row.SQL, "__legacy_budget_backup", "table %s has FK referencing a legacy backup table: %s", row.Name, row.SQL) assert.NotContains(t, row.SQL, "__new", "table %s has FK referencing a temp table: %s", row.Name, row.SQL) } } func strPtr(s string) *string { return &s }