netbird/management/server/migration/migration_test.go

package migration_test

import (
	"context"
	"encoding/gob"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
	"github.com/netbirdio/netbird/management/server/migration"
	nbpeer "github.com/netbirdio/netbird/management/server/peer"
	"github.com/netbirdio/netbird/management/server/testutil"
	"github.com/netbirdio/netbird/management/server/types"
	"github.com/netbirdio/netbird/route"
)

func setupDatabase(t *testing.T) *gorm.DB {
	t.Helper()

	var db *gorm.DB
	var err error
	var dsn string
	var cleanup func()
	switch os.Getenv("NETBIRD_STORE_ENGINE") {
	case "mysql":
		cleanup, dsn, err = testutil.CreateMysqlTestContainer()
		if err != nil {
			t.Fatalf("Failed to create MySQL test container: %v", err)
		}

		if dsn == "" {
			t.Fatal("MySQL connection string is empty, ensure the test container is running")
		}

		db, err = gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), &gorm.Config{})
	case "postgres":
		cleanup, dsn, err = testutil.CreatePostgresTestContainer()
		if err != nil {
			t.Fatalf("Failed to create PostgreSQL test container: %v", err)
		}

		if dsn == "" {
			t.Fatalf("PostgreSQL connection string is empty, ensure the test container is running")
		}

		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
	case "sqlite":
		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	default:
		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	}
	if cleanup != nil {
		t.Cleanup(cleanup)
	}

	require.NoError(t, err, "Failed to open database")
	return db
}

func TestMigrateFieldFromGobToJSON_EmptyDB(t *testing.T) {
	db := setupDatabase(t)
	err := migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail for an empty database")
}

func TestMigrateFieldFromGobToJSON_WithGobData(t *testing.T) {
	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &route.Route{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	_, ipnet, err := net.ParseCIDR("10.0.0.0/24")
	require.NoError(t, err, "Failed to parse CIDR")

	type network struct {
		types.Network
		Net net.IPNet `gorm:"serializer:gob"`
	}

	type account struct {
		types.Account
		Network *network `gorm:"embedded;embeddedPrefix:network_"`
	}

	err = db.Save(&account{Account: types.Account{Id: "123"}, Network: &network{Net: *ipnet}}).Error
	require.NoError(t, err, "Failed to insert Gob data")

	var gobStr string
	err = db.Model(&types.Account{}).Select("network_net").First(&gobStr).Error
	assert.NoError(t, err, "Failed to fetch Gob data")

	err = gob.NewDecoder(strings.NewReader(gobStr)).Decode(&ipnet)
	require.NoError(t, err, "Failed to decode Gob data")

	err = migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail with Gob data")

	var jsonStr string
	db.Model(&types.Account{}).Select("network_net").First(&jsonStr)
	assert.JSONEq(t, `{"IP":"10.0.0.0","Mask":"////AA=="}`, jsonStr, "Data should be migrated")
}

func TestMigrateFieldFromGobToJSON_WithJSONData(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &route.Route{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	_, ipnet, err := net.ParseCIDR("10.0.0.0/24")
	require.NoError(t, err, "Failed to parse CIDR")

	err = db.Save(&types.Account{Network: &types.Network{Net: *ipnet}}).Error
	require.NoError(t, err, "Failed to insert JSON data")

	err = migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail with JSON data")

	var jsonStr string
	db.Model(&types.Account{}).Select("network_net").First(&jsonStr)
	assert.JSONEq(t, `{"IP":"10.0.0.0","Mask":"////AA=="}`, jsonStr, "Data should be unchanged")
}

func TestMigrateNetIPFieldFromBlobToJSON_EmptyDB(t *testing.T) {
	db := setupDatabase(t)
	err := migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "ip", "idx_peers_account_id_ip")
	require.NoError(t, err, "Migration should not fail for an empty database")
}

func TestMigrateNetIPFieldFromBlobToJSON_WithBlobData(t *testing.T) {
	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	type location struct {
		nbpeer.Location
		ConnectionIP net.IP
	}

	type peer struct {
		nbpeer.Peer
		Location location `gorm:"embedded;embeddedPrefix:location_"`
	}

	type account struct {
		types.Account
		Peers []peer `gorm:"foreignKey:AccountID;references:id"`
	}

	a := &account{
		Account: types.Account{Id: "123"},
	}

	err = db.Save(a).Error
	require.NoError(t, err, "Failed to insert account")

	a.Peers = []peer{
		{Location: location{ConnectionIP: net.IP{10, 0, 0, 1}}},
	}

	err = db.Save(a).Error
	require.NoError(t, err, "Failed to insert blob data")

	var blobValue string
	err = db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&blobValue).Error
	assert.NoError(t, err, "Failed to fetch blob data")

	err = migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "location_connection_ip", "")
	require.NoError(t, err, "Migration should not fail with net.IP blob data")

	var jsonStr string
	db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&jsonStr)
	assert.JSONEq(t, `"10.0.0.1"`, jsonStr, "Data should be migrated")
}

func TestMigrateNetIPFieldFromBlobToJSON_WithJSONData(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	account := &types.Account{
		Id: "1234",
	}

	err = db.Save(account).Error
	require.NoError(t, err, "Failed to insert account")

	account.PeersG = []nbpeer.Peer{
		{
			AccountID: "1234",
			Location:  nbpeer.Location{ConnectionIP: net.IP{10, 0, 0, 1}},
			Status:    &nbpeer.PeerStatus{LastSeen: time.Now()},
		},
	}

	err = db.Save(account).Error
	require.NoError(t, err, "Failed to insert JSON data")

	err = migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "location_connection_ip", "")
	require.NoError(t, err, "Migration should not fail with net.IP JSON data")

	var jsonStr string
	db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&jsonStr)
	assert.JSONEq(t, `"10.0.0.1"`, jsonStr, "Data should be unchanged")
}

func TestMigrateSetupKeyToHashedSetupKey_ForPlainKey(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "EEFDAB47-C1A5-4472-8C05-71DE9A1E8382",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "EEFDA****", key.KeySecret, "Key should be secret")
	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case1(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		KeySecret: "EEFDA****",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "EEFDA****", key.KeySecret, "Key should be secret")
	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case2(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestDropIndex(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	exist := db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
	assert.True(t, exist, "Should have the index")

	err = migration.DropIndex[types.SetupKey](context.Background(), db, "idx_setup_keys_account_id")
	require.NoError(t, err, "Migration should not fail to remove index")

	exist = db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
	assert.False(t, exist, "Should not have the index")
}

func TestCreateIndex(t *testing.T) {
	db := setupDatabase(t)
	err := db.AutoMigrate(&nbpeer.Peer{})
	assert.NoError(t, err, "Failed to auto-migrate tables")

	indexName := "idx_account_ip"

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Migration should not fail to create index")

	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")
}

func TestCreateIndexIfExists(t *testing.T) {
	db := setupDatabase(t)
	err := db.AutoMigrate(&nbpeer.Peer{})
	assert.NoError(t, err, "Failed to auto-migrate tables")

	indexName := "idx_account_ip"

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Migration should not fail to create index")

	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Create index should not fail if index exists")

	exist = db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")
}

type testPeer struct {
	ID                  string `gorm:"primaryKey"`
	Key                 string `gorm:"index"`
	PeerStatusLastSeen  time.Time
	PeerStatusConnected bool
}

func (testPeer) TableName() string {
	return "peers"
}

func setupPeerTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	db := setupDatabase(t)
	_ = db.Migrator().DropTable(&testPeer{})
	err := db.AutoMigrate(&testPeer{})
	require.NoError(t, err, "Failed to auto-migrate tables")
	return db
}

func TestRemoveDuplicatePeerKeys_NoDuplicates(t *testing.T) {
	db := setupPeerTestDB(t)

	now := time.Now()
	peers := []testPeer{
		{ID: "peer1", Key: "key1", PeerStatusLastSeen: now},
		{ID: "peer2", Key: "key2", PeerStatusLastSeen: now},
		{ID: "peer3", Key: "key3", PeerStatusLastSeen: now},
	}

	for _, p := range peers {
		err := db.Create(&p).Error
		require.NoError(t, err)
	}

	err := migration.RemoveDuplicatePeerKeys(context.Background(), db)
	require.NoError(t, err)

	var count int64
	db.Model(&testPeer{}).Count(&count)
	assert.Equal(t, int64(len(peers)), count, "All peers should remain when no duplicates")
}

func TestRemoveDuplicatePeerKeys_WithDuplicates(t *testing.T) {
	db := setupPeerTestDB(t)

	now := time.Now()
	peers := []testPeer{
		{ID: "peer1", Key: "key1", PeerStatusLastSeen: now.Add(-2 * time.Hour)},
		{ID: "peer2", Key: "key1", PeerStatusLastSeen: now.Add(-1 * time.Hour)},
		{ID: "peer3", Key: "key1", PeerStatusLastSeen: now},
		{ID: "peer4", Key: "key2", PeerStatusLastSeen: now},
		{ID: "peer5", Key: "key3", PeerStatusLastSeen: now.Add(-1 * time.Hour)},
		{ID: "peer6", Key: "key3", PeerStatusLastSeen: now},
	}

	for _, p := range peers {
		err := db.Create(&p).Error
		require.NoError(t, err)
	}

	err := migration.RemoveDuplicatePeerKeys(context.Background(), db)
	require.NoError(t, err)

	var count int64
	db.Model(&testPeer{}).Count(&count)
	assert.Equal(t, int64(3), count, "Should have 3 peers after removing duplicates")

	var remainingPeers []testPeer
	err = db.Find(&remainingPeers).Error
	require.NoError(t, err)

	remainingIDs := make(map[string]bool)
	for _, p := range remainingPeers {
		remainingIDs[p.ID] = true
	}

	assert.True(t, remainingIDs["peer3"], "peer3 should remain (most recent for key1)")
	assert.True(t, remainingIDs["peer4"], "peer4 should remain (only peer for key2)")
	assert.True(t, remainingIDs["peer6"], "peer6 should remain (most recent for key3)")

	assert.False(t, remainingIDs["peer1"], "peer1 should be deleted (older duplicate)")
	assert.False(t, remainingIDs["peer2"], "peer2 should be deleted (older duplicate)")
	assert.False(t, remainingIDs["peer5"], "peer5 should be deleted (older duplicate)")
}

func TestRemoveDuplicatePeerKeys_EmptyTable(t *testing.T) {
	db := setupPeerTestDB(t)

	err := migration.RemoveDuplicatePeerKeys(context.Background(), db)
	require.NoError(t, err, "Should not fail on empty table")
}

func TestRemoveDuplicatePeerKeys_NoTable(t *testing.T) {
	db := setupDatabase(t)
	_ = db.Migrator().DropTable(&testPeer{})

	err := migration.RemoveDuplicatePeerKeys(context.Background(), db)
	require.NoError(t, err, "Should not fail when table does not exist")
}

type testParent struct {
	ID string `gorm:"primaryKey"`
}

func (testParent) TableName() string {
	return "test_parents"
}

type testChild struct {
	ID       string `gorm:"primaryKey"`
	ParentID string
}

func (testChild) TableName() string {
	return "test_children"
}

type testChildWithFK struct {
	ID       string      `gorm:"primaryKey"`
	ParentID string      `gorm:"index"`
	Parent   *testParent `gorm:"foreignKey:ParentID"`
}

func (testChildWithFK) TableName() string {
	return "test_children"
}

func setupOrphanTestDB(t *testing.T, models ...any) *gorm.DB {
	t.Helper()
	db := setupDatabase(t)
	for _, m := range models {
		_ = db.Migrator().DropTable(m)
	}
	err := db.AutoMigrate(models...)
	require.NoError(t, err, "Failed to auto-migrate tables")
	return db
}

func TestCleanupOrphanedResources_NoChildTable(t *testing.T) {
	db := setupDatabase(t)
	_ = db.Migrator().DropTable(&testChild{})
	_ = db.Migrator().DropTable(&testParent{})

	err := migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err, "Should not fail when child table does not exist")
}

func TestCleanupOrphanedResources_NoParentTable(t *testing.T) {
	db := setupDatabase(t)
	_ = db.Migrator().DropTable(&testParent{})
	_ = db.Migrator().DropTable(&testChild{})

	err := db.AutoMigrate(&testChild{})
	require.NoError(t, err)

	err = migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err, "Should not fail when parent table does not exist")
}

func TestCleanupOrphanedResources_EmptyTables(t *testing.T) {
	db := setupOrphanTestDB(t, &testParent{}, &testChild{})

	err := migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err, "Should not fail on empty tables")

	var count int64
	db.Model(&testChild{}).Count(&count)
	assert.Equal(t, int64(0), count)
}

func TestCleanupOrphanedResources_NoOrphans(t *testing.T) {
	db := setupOrphanTestDB(t, &testParent{}, &testChild{})

	require.NoError(t, db.Create(&testParent{ID: "p1"}).Error)
	require.NoError(t, db.Create(&testParent{ID: "p2"}).Error)
	require.NoError(t, db.Create(&testChild{ID: "c1", ParentID: "p1"}).Error)
	require.NoError(t, db.Create(&testChild{ID: "c2", ParentID: "p2"}).Error)

	err := migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err)

	var count int64
	db.Model(&testChild{}).Count(&count)
	assert.Equal(t, int64(2), count, "All children should remain when no orphans")
}

func TestCleanupOrphanedResources_AllOrphans(t *testing.T) {
	db := setupOrphanTestDB(t, &testParent{}, &testChild{})

	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c1", "gone1").Error)
	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c2", "gone2").Error)
	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c3", "gone3").Error)

	err := migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err)

	var count int64
	db.Model(&testChild{}).Count(&count)
	assert.Equal(t, int64(0), count, "All orphaned children should be deleted")
}

func TestCleanupOrphanedResources_MixedValidAndOrphaned(t *testing.T) {
	db := setupOrphanTestDB(t, &testParent{}, &testChild{})

	require.NoError(t, db.Create(&testParent{ID: "p1"}).Error)
	require.NoError(t, db.Create(&testParent{ID: "p2"}).Error)

	require.NoError(t, db.Create(&testChild{ID: "c1", ParentID: "p1"}).Error)
	require.NoError(t, db.Create(&testChild{ID: "c2", ParentID: "p2"}).Error)
	require.NoError(t, db.Create(&testChild{ID: "c3", ParentID: "p1"}).Error)

	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c4", "gone1").Error)
	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c5", "gone2").Error)

	err := migration.CleanupOrphanedResources[testChild, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err)

	var remaining []testChild
	require.NoError(t, db.Order("id").Find(&remaining).Error)

	assert.Len(t, remaining, 3, "Only valid children should remain")
	assert.Equal(t, "c1", remaining[0].ID)
	assert.Equal(t, "c2", remaining[1].ID)
	assert.Equal(t, "c3", remaining[2].ID)
}

func TestCleanupOrphanedResources_Idempotent(t *testing.T) {
	db := setupOrphanTestDB(t, &testParent{}, &testChild{})

	require.NoError(t, db.Create(&testParent{ID: "p1"}).Error)
	require.NoError(t, db.Create(&testChild{ID: "c1", ParentID: "p1"}).Error)
	require.NoError(t, db.Exec("INSERT INTO test_children (id, parent_id) VALUES (?, ?)", "c2", "gone").Error)

	ctx := context.Background()

	err := migration.CleanupOrphanedResources[testChild, testParent](ctx, db, "parent_id")
	require.NoError(t, err)

	var count int64
	db.Model(&testChild{}).Count(&count)
	assert.Equal(t, int64(1), count)

	err = migration.CleanupOrphanedResources[testChild, testParent](ctx, db, "parent_id")
	require.NoError(t, err)

	db.Model(&testChild{}).Count(&count)
	assert.Equal(t, int64(1), count, "Count should remain the same after second run")
}

func TestCleanupOrphanedResources_SkipsWhenForeignKeyExists(t *testing.T) {
	engine := os.Getenv("NETBIRD_STORE_ENGINE")
	if engine != "postgres" && engine != "mysql" {
		t.Skip("FK constraint early-exit test requires postgres or mysql")
	}

	db := setupDatabase(t)
	_ = db.Migrator().DropTable(&testChildWithFK{})
	_ = db.Migrator().DropTable(&testParent{})

	err := db.AutoMigrate(&testParent{}, &testChildWithFK{})
	require.NoError(t, err)

	require.NoError(t, db.Create(&testParent{ID: "p1"}).Error)
	require.NoError(t, db.Create(&testParent{ID: "p2"}).Error)
	require.NoError(t, db.Create(&testChildWithFK{ID: "c1", ParentID: "p1"}).Error)
	require.NoError(t, db.Create(&testChildWithFK{ID: "c2", ParentID: "p2"}).Error)

	switch engine {
	case "postgres":
		require.NoError(t, db.Exec("ALTER TABLE test_children DROP CONSTRAINT fk_test_children_parent").Error)
		require.NoError(t, db.Exec("DELETE FROM test_parents WHERE id = ?", "p2").Error)
		require.NoError(t, db.Exec(
			"ALTER TABLE test_children ADD CONSTRAINT fk_test_children_parent "+
				"FOREIGN KEY (parent_id) REFERENCES test_parents(id) NOT VALID",
		).Error)
	case "mysql":
		require.NoError(t, db.Exec("SET FOREIGN_KEY_CHECKS = 0").Error)
		require.NoError(t, db.Exec("ALTER TABLE test_children DROP FOREIGN KEY fk_test_children_parent").Error)
		require.NoError(t, db.Exec("DELETE FROM test_parents WHERE id = ?", "p2").Error)
		require.NoError(t, db.Exec(
			"ALTER TABLE test_children ADD CONSTRAINT fk_test_children_parent "+
				"FOREIGN KEY (parent_id) REFERENCES test_parents(id)",
		).Error)
		require.NoError(t, db.Exec("SET FOREIGN_KEY_CHECKS = 1").Error)
	}

	err = migration.CleanupOrphanedResources[testChildWithFK, testParent](context.Background(), db, "parent_id")
	require.NoError(t, err)

	var count int64
	db.Model(&testChildWithFK{}).Count(&count)
	assert.Equal(t, int64(2), count, "Both rows should survive — migration must skip when FK constraint exists")
}

// legacyCostRow is the pre-breakdown shape of the usage table: cost was stored
// as a total plus a cache portion, with no per-bucket columns. Used to build a
// realistic pre-upgrade table for the fold migration to run against.
type legacyCostRow struct {
	ID           string `gorm:"primaryKey"`
	AccountID    string
	Model        string
	CostUSD      float64
	CacheCostUSD float64
}

func (legacyCostRow) TableName() string { return "agent_network_request_usage" }

// TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost covers the upgrade
// path: a table written under the old schema must come out with its per-row
// total and cache cost unchanged, because dropping cost_usd without folding it
// forward would silently zero every historical row's spend.
func TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	// setupDatabase hands back a process-shared database, so start from a clean
	// table rather than inheriting rows from another test.
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))

	require.NoError(t, db.AutoMigrate(&legacyCostRow{}), "legacy table must be created")
	require.NoError(t, db.Create(&legacyCostRow{
		ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6", CostUSD: 0.0123, CacheCostUSD: 0.0029,
	}).Error)
	require.NoError(t, db.Create(&legacyCostRow{
		ID: "u2", AccountID: "acct-1", Model: "gpt-4o", CostUSD: 0.5, CacheCostUSD: 0,
	}).Error)
	// A zero-cost row (denied / unpriced request) must stay zero, not be touched.
	require.NoError(t, db.Create(&legacyCostRow{ID: "u3", AccountID: "acct-1", Model: "gw/unpriced"}).Error)

	// AutoMigrate adds the per-bucket columns alongside the legacy ones, exactly
	// as a real upgrade does before the post-auto migrations run.
	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}), "new columns must be added")

	require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db))

	assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cost_usd"),
		"legacy cost_usd column must be dropped once folded")
	assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cache_cost_usd"),
		"legacy cache_cost_usd column must be dropped once folded")

	var rows []*agentNetworkTypes.AgentNetworkUsage
	require.NoError(t, db.Order("id").Find(&rows).Error)
	require.Len(t, rows, 3)

	// u1: total and cache portion both preserved; the read/write and
	// input/output splits are unknowable for a legacy row, so the cache total
	// lands on cached_input and the remainder on input.
	assert.InDelta(t, 0.0123, rows[0].TotalCostUSD(), 1e-9, "historical total must survive the fold")
	assert.InDelta(t, 0.0029, rows[0].CacheCostUSD(), 1e-9, "historical cache cost must survive the fold")
	assert.InDelta(t, 0.0094, rows[0].InputCostUSD, 1e-9, "non-cache remainder lands on input")
	assert.InDelta(t, 0.0029, rows[0].CachedInputCostUSD, 1e-9, "legacy cache total lands on cached input")
	assert.Zero(t, rows[0].CacheCreationCostUSD, "legacy rows carry no read/write split to recover")
	assert.Zero(t, rows[0].OutputCostUSD, "legacy rows carry no input/output split to recover")

	// u2: no cache spend — the whole total is the non-cache remainder.
	assert.InDelta(t, 0.5, rows[1].TotalCostUSD(), 1e-9, "cache-free historical total must survive")
	assert.Zero(t, rows[1].CacheCostUSD(), "a cache-free row must stay cache-free")

	// u3: zero stays zero rather than being rewritten.
	assert.Zero(t, rows[2].TotalCostUSD(), "an unpriced row must remain unpriced")
}

// TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated proves the migration is
// safe to re-run: with no legacy column present it is a no-op that leaves a
// true four-way split untouched.
func TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))

	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}))
	// Timestamp must be set explicitly: a zero time.Time serialises as
	// '0000-00-00 00:00:00', which MySQL rejects under strict mode.
	require.NoError(t, db.Create(&agentNetworkTypes.AgentNetworkUsage{
		ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6",
		Timestamp:    time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC),
		InputCostUSD: 0.001, CachedInputCostUSD: 0.002, CacheCreationCostUSD: 0.003, OutputCostUSD: 0.004,
	}).Error)

	require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db),
		"running against an already-migrated table must be a no-op, not an error")

	var row agentNetworkTypes.AgentNetworkUsage
	require.NoError(t, db.First(&row, "id = ?", "u1").Error)
	assert.InDelta(t, 0.001, row.InputCostUSD, 1e-9, "a true split must not be rewritten")
	assert.InDelta(t, 0.002, row.CachedInputCostUSD, 1e-9)
	assert.InDelta(t, 0.003, row.CacheCreationCostUSD, 1e-9)
	assert.InDelta(t, 0.004, row.OutputCostUSD, 1e-9)
	assert.InDelta(t, 0.01, row.TotalCostUSD(), 1e-9, "derived total sums the four buckets")
}

// legacyAgentNetworkSettings is the pre-reshape schema: identity carried as
// (cluster, subdomain) instead of (domain, proxy_address).
type legacyAgentNetworkSettings struct {
	AccountID           string `gorm:"primaryKey"`
	Cluster             string
	Subdomain           string
	EnableLogCollection bool
}

func (legacyAgentNetworkSettings) TableName() string { return "agent_network_settings" }

// TestMigrateAgentNetworkSettingsToDomain_BackfillsAndDropsLegacyColumns pins
// the reshape: domain becomes `<subdomain>.<cluster>`, proxy_address becomes
// the cluster, the legacy columns are dropped, and non-identity fields ride
// through untouched.
func TestMigrateAgentNetworkSettingsToDomain_BackfillsAndDropsLegacyColumns(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&legacyAgentNetworkSettings{}))
	require.NoError(t, db.AutoMigrate(&legacyAgentNetworkSettings{}))
	// The cluster is spelled the way the legacy bootstrap kept it: as the
	// caller typed it, trimmed but never folded. The subdomain was always
	// server-assigned lowercase.
	require.NoError(t, db.Create(&legacyAgentNetworkSettings{
		AccountID: "acct-1", Cluster: "EU.Proxy.NetBird.io", Subdomain: "violet", EnableLogCollection: true,
	}).Error)
	require.NoError(t, db.Create(&legacyAgentNetworkSettings{
		AccountID: "acct-2", Cluster: "us.proxy.netbird.io", Subdomain: "violet",
	}).Error)

	require.NoError(t, migration.MigrateAgentNetworkSettingsToDomain(ctx, db))
	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.Settings{}),
		"AutoMigrate must create the domain unique index over the backfilled values")

	var one, two agentNetworkTypes.Settings
	require.NoError(t, db.First(&one, "account_id = ?", "acct-1").Error)
	assert.Equal(t, "violet.eu.proxy.netbird.io", one.Domain,
		"domain must combine subdomain and cluster, folded to the canonical lowercase every reader compares against")
	assert.Equal(t, "eu.proxy.netbird.io", one.ProxyAddress,
		"proxy address must carry the cluster in canonical lowercase, matching what proxies register under")
	assert.True(t, one.EnableLogCollection, "non-identity fields must ride through")
	require.NoError(t, db.First(&two, "account_id = ?", "acct-2").Error)
	assert.Equal(t, "violet.us.proxy.netbird.io", two.Domain,
		"duplicate labels on different clusters are distinct hostnames and must both survive")

	migrator := db.Migrator()
	assert.False(t, migrator.HasColumn(&legacyAgentNetworkSettings{}, "cluster"), "legacy cluster column must be dropped")
	assert.False(t, migrator.HasColumn(&legacyAgentNetworkSettings{}, "subdomain"), "legacy subdomain column must be dropped")
}

// TestMigrateAgentNetworkSettingsToDomain_SkipsAlreadyMigrated proves the
// migration is safe to re-run: with no legacy column present it is a no-op.
func TestMigrateAgentNetworkSettingsToDomain_SkipsAlreadyMigrated(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.Create(&agentNetworkTypes.Settings{
		AccountID: "acct-1", Domain: "gw.example.com", ProxyAddress: "gw.example.com",
	}).Error)

	require.NoError(t, migration.MigrateAgentNetworkSettingsToDomain(ctx, db),
		"running against an already-migrated table must be a no-op, not an error")

	var row agentNetworkTypes.Settings
	require.NoError(t, db.First(&row, "account_id = ?", "acct-1").Error)
	assert.Equal(t, "gw.example.com", row.Domain, "migrated rows must be untouched")
}

// TestMigrateAgentNetworkSettingsToDomain_FailsOnUnmigratableRow pins the
// loud-failure contract: a legacy row missing its identity halves cannot be
// given an endpoint, and silently leaving an empty domain would collide with
// the unique index confusingly later.
func TestMigrateAgentNetworkSettingsToDomain_FailsOnUnmigratableRow(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&legacyAgentNetworkSettings{}))
	require.NoError(t, db.AutoMigrate(&legacyAgentNetworkSettings{}))
	require.NoError(t, db.Create(&legacyAgentNetworkSettings{
		AccountID: "acct-broken", Cluster: "", Subdomain: "",
	}).Error)

	err := migration.MigrateAgentNetworkSettingsToDomain(ctx, db)
	require.Error(t, err, "a row with no identity to derive an endpoint from must fail the migration")
	assert.Contains(t, err.Error(), "resolve them manually", "the error must tell the operator what to do")
}

// partialAgentNetworkSettings models the one non-atomic state a MySQL run can
// be interrupted in: DDL auto-commits there, so a crash between the two legacy
// column drops leaves subdomain behind while cluster (and the completed
// backfill) are already committed.
type partialAgentNetworkSettings struct {
	AccountID    string `gorm:"primaryKey"`
	Subdomain    string
	Domain       string `gorm:"type:varchar(255)"`
	ProxyAddress string `gorm:"type:varchar(255)"`
}

func (partialAgentNetworkSettings) TableName() string { return "agent_network_settings" }

// TestMigrateAgentNetworkSettingsToDomain_ResumesAfterPartialDrop pins MySQL
// resumability: a rerun over the interrupted state must remove the leftover
// subdomain column without re-running the backfill (the cluster column that
// feeds it is gone) and without touching the migrated values.
func TestMigrateAgentNetworkSettingsToDomain_ResumesAfterPartialDrop(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&partialAgentNetworkSettings{}))
	require.NoError(t, db.AutoMigrate(&partialAgentNetworkSettings{}))
	require.NoError(t, db.Create(&partialAgentNetworkSettings{
		AccountID: "acct-1", Subdomain: "violet",
		Domain: "violet.eu.proxy.netbird.io", ProxyAddress: "eu.proxy.netbird.io",
	}).Error)

	require.NoError(t, migration.MigrateAgentNetworkSettingsToDomain(ctx, db),
		"a rerun over a partially-dropped schema must resume, not error")

	migrator := db.Migrator()
	assert.False(t, migrator.HasColumn(&partialAgentNetworkSettings{}, "subdomain"),
		"the leftover legacy column must be dropped on resume")

	var row agentNetworkTypes.Settings
	require.NoError(t, db.First(&row, "account_id = ?", "acct-1").Error)
	assert.Equal(t, "violet.eu.proxy.netbird.io", row.Domain, "migrated values must be untouched")
	assert.Equal(t, "eu.proxy.netbird.io", row.ProxyAddress, "migrated values must be untouched")
}

// TestNormalizeAgentNetworkSettingsIdentity_LowercasesReshapedRows covers rows
// a released reshape already copied verbatim: capitals kept from the legacy
// cluster spelling are folded in place, canonical rows are left alone, and
// non-identity fields ride through.
func TestNormalizeAgentNetworkSettingsIdentity_LowercasesReshapedRows(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.Create(&agentNetworkTypes.Settings{
		AccountID: "acct-legacy", Domain: "Violet.EU.Proxy.NetBird.io", ProxyAddress: "EU.Proxy.NetBird.io", EnableLogCollection: true,
	}).Error)
	require.NoError(t, db.Create(&agentNetworkTypes.Settings{
		AccountID: "acct-canonical", Domain: "amber.us.proxy.netbird.io", ProxyAddress: "us.proxy.netbird.io",
	}).Error)

	require.NoError(t, migration.NormalizeAgentNetworkSettingsIdentity(ctx, db))

	var legacy, canonical agentNetworkTypes.Settings
	require.NoError(t, db.First(&legacy, "account_id = ?", "acct-legacy").Error)
	assert.Equal(t, "violet.eu.proxy.netbird.io", legacy.Domain, "a mixed-case endpoint must be folded where it is stored")
	assert.Equal(t, "eu.proxy.netbird.io", legacy.ProxyAddress, "a mixed-case pin must be folded so exact lookups find it")
	assert.True(t, legacy.EnableLogCollection, "non-identity fields must ride through")
	require.NoError(t, db.First(&canonical, "account_id = ?", "acct-canonical").Error)
	assert.Equal(t, "amber.us.proxy.netbird.io", canonical.Domain, "a canonical row must be left as it is")
	assert.Equal(t, "us.proxy.netbird.io", canonical.ProxyAddress)

	require.NoError(t, migration.NormalizeAgentNetworkSettingsIdentity(ctx, db),
		"a second run over a normalised table must be a no-op, not an error")
}

// TestNormalizeAgentNetworkSettingsIdentity_SkipsMissingTable pins that a
// store which never had agent network settings is left untouched.
func TestNormalizeAgentNetworkSettingsIdentity_SkipsMissingTable(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.Settings{}))

	require.NoError(t, migration.NormalizeAgentNetworkSettingsIdentity(ctx, db),
		"no table must be a no-op, not an error")
	assert.False(t, db.Migrator().HasTable(&agentNetworkTypes.Settings{}), "the normaliser must not create the table")
}

// TestNormalizeAgentNetworkSettingsIdentity_RefusesCaseOnlyCollision pins the
// loud failure: two rows that would fold onto one endpoint stop the migration
// with the hostname named, and neither row is touched, rather than letting the
// unique index refuse the fold with a driver message that names no row.
func TestNormalizeAgentNetworkSettingsIdentity_RefusesCaseOnlyCollision(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	if db.Name() == "mysql" {
		t.Skip("MySQL's default collation refuses two rows differing only by case at insert; the collision cannot exist there")
	}
	require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.AutoMigrate(&agentNetworkTypes.Settings{}))
	require.NoError(t, db.Create(&agentNetworkTypes.Settings{
		AccountID: "acct-1", Domain: "Violet.eu.proxy.netbird.io", ProxyAddress: "eu.proxy.netbird.io",
	}).Error)
	require.NoError(t, db.Create(&agentNetworkTypes.Settings{
		AccountID: "acct-2", Domain: "violet.eu.proxy.netbird.io", ProxyAddress: "eu.proxy.netbird.io",
	}).Error)

	err := migration.NormalizeAgentNetworkSettingsIdentity(ctx, db)
	require.Error(t, err, "two rows folding onto one endpoint must stop the migration")
	assert.Contains(t, err.Error(), "violet.eu.proxy.netbird.io", "the failure must name the colliding hostname")

	var one agentNetworkTypes.Settings
	require.NoError(t, db.First(&one, "account_id = ?", "acct-1").Error)
	assert.Equal(t, "Violet.eu.proxy.netbird.io", one.Domain, "a refused normalisation must leave every row as it was")
}

// TestMigrateAgentNetworkSettingsToDomain_RefusesCaseOnlyCollision pins the
// same loud failure on the reshape: legacy rows whose identities differ only
// by case would fold onto one endpoint, and the reshape must say so rather
// than leave AutoMigrate to fail on the unique index.
func TestMigrateAgentNetworkSettingsToDomain_RefusesCaseOnlyCollision(t *testing.T) {
	ctx := context.Background()
	db := setupDatabase(t)
	require.NoError(t, db.Migrator().DropTable(&legacyAgentNetworkSettings{}))
	require.NoError(t, db.AutoMigrate(&legacyAgentNetworkSettings{}))
	require.NoError(t, db.Create(&legacyAgentNetworkSettings{
		AccountID: "acct-1", Cluster: "EU.proxy.netbird.io", Subdomain: "violet",
	}).Error)
	require.NoError(t, db.Create(&legacyAgentNetworkSettings{
		AccountID: "acct-2", Cluster: "eu.proxy.netbird.io", Subdomain: "violet",
	}).Error)

	err := migration.MigrateAgentNetworkSettingsToDomain(ctx, db)
	require.Error(t, err, "legacy rows folding onto one endpoint must stop the reshape")
	assert.Contains(t, err.Error(), "violet.eu.proxy.netbird.io", "the failure must name the colliding hostname")
}