sub2api-cn-relay-manager/internal/store/sqlite/db_test.go

package sqlite

import (
	"context"
	"database/sql"
	"errors"
	"path/filepath"
	"testing"
)

func TestOpenClose(t *testing.T) {
	store := openTestDB(t)
	if store == nil {
		t.Fatal("Open() returned nil")
	}
}

func TestOpenPinsSingleSQLiteConnection(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "single-conn.db")
	store, err := Open(context.Background(), "file:"+filepath.ToSlash(dbPath)+"?_busy_timeout=5000")
	if err != nil {
		t.Fatalf("Open() error = %v", err)
	}
	defer func() {
		if err := store.Close(); err != nil {
			t.Fatalf("Close() error = %v", err)
		}
	}()

	stats := store.SQLDB().Stats()
	if stats.MaxOpenConnections != 1 {
		t.Fatalf("MaxOpenConnections = %d, want 1", stats.MaxOpenConnections)
	}
}

func TestOpenInvalidDSN(t *testing.T) {
	_, err := Open(context.Background(), "file:/nonexistent/dir/test.db?_pragma=foreign_keys(0)")
	if err == nil {
		t.Fatal("Open() with invalid dsn error = nil, want error")
	}
}

func TestWithTxCommit(t *testing.T) {
	store := openTestDB(t)

	err := store.WithTx(context.Background(), func(q *Queries) error {
		_, err := q.Hosts.Create(context.Background(), Host{
			HostID: "tx-host", BaseURL: "https://tx.com", HostVersion: "1.0",
		})
		return err
	})
	if err != nil {
		t.Fatalf("WithTx() error = %v", err)
	}

	host, err := store.Hosts().GetByHostID(context.Background(), "tx-host")
	if err != nil {
		t.Fatalf("GetByHostID() after tx = %v", err)
	}
	if host.HostID != "tx-host" {
		t.Fatalf("host = %+v, want tx-host", host)
	}
}

func TestWithTxRollbackOnError(t *testing.T) {
	store := openTestDB(t)

	err := store.WithTx(context.Background(), func(q *Queries) error {
		q.Hosts.Create(context.Background(), Host{
			HostID: "rollback-host", BaseURL: "https://r.com", HostVersion: "1.0",
		})
		return errors.New("rollback")
	})
	if err == nil {
		t.Fatal("WithTx() error = nil, want rollback error")
	}

	_, err = store.Hosts().GetByHostID(context.Background(), "rollback-host")
	if !errors.Is(err, sql.ErrNoRows) {
		t.Fatalf("GetByHostID() after rollback error = %v, want sql.ErrNoRows", err)
	}
}

func TestTableExists(t *testing.T) {
	store := openTestDB(t)
	db := store.SQLDB()

	found, err := tableExists(context.Background(), db, "hosts")
	if err != nil {
		t.Fatalf("tableExists('hosts') error = %v", err)
	}
	if !found {
		t.Fatal("tableExists('hosts') = false, want true")
	}

	found, err = tableExists(context.Background(), db, "nonexistent")
	if err != nil {
		t.Fatalf("tableExists('nonexistent') error = %v", err)
	}
	if found {
		t.Fatal("tableExists('nonexistent') = true, want false")
	}
}

func TestOpenAppliesLogicalRoutingTables(t *testing.T) {
	store := openTestDB(t)
	db := store.SQLDB()

	for _, table := range []string{
		"logical_groups",
		"logical_group_models",
		"logical_group_routes",
		"logical_group_route_models",
		"route_decision_logs",
		"route_failover_events",
		"route_sticky_audit",
		"provider_accounts",
	} {
		found, err := tableExists(context.Background(), db, table)
		if err != nil {
			t.Fatalf("tableExists(%q) error = %v", table, err)
		}
		if !found {
			t.Fatalf("tableExists(%q) = false, want true", table)
		}
	}
}

func TestDetectLegacy0001Schema(t *testing.T) {
	store := openTestDB(t)
	db := store.SQLDB()

	tx, err := db.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatalf("BeginTx error = %v", err)
	}
	defer tx.Rollback()

	// After migration all three host/packs/providers tables exist,
	// so detectLegacy0001Schema reports complete=true.
	complete, partial, err := detectLegacy0001Schema(context.Background(), tx)
	if err != nil {
		t.Fatalf("detectLegacy0001Schema() error = %v", err)
	}
	if !complete {
		t.Fatalf("detectLegacy0001Schema() = (complete=%v, partial=%v), want (true, false)", complete, partial)
	}
	if partial {
		t.Fatal("partial should be false when all 3 tables exist")
	}
}

func TestWithForeignKeysEnabled(t *testing.T) {
	if got := withForeignKeysEnabled("file:test.db"); got != "file:test.db?_pragma=foreign_keys(1)" {
		t.Fatalf("withForeignKeysEnabled no query = %q", got)
	}
	if got := withForeignKeysEnabled("file:test.db?a=1"); got != "file:test.db?a=1&_pragma=foreign_keys(1)" {
		t.Fatalf("withForeignKeysEnabled with query = %q", got)
	}
}

func TestSQLDB(t *testing.T) {
	store := openTestDB(t)
	db := store.SQLDB()
	if db == nil {
		t.Fatal("SQLDB() returned nil")
	}
	if err := db.PingContext(context.Background()); err != nil {
		t.Fatalf("Ping() error = %v", err)
	}
}

func TestMigrationFileNames(t *testing.T) {
	names, err := migrationFileNames()
	if err != nil {
		t.Fatalf("migrationFileNames() error = %v", err)
	}
	if len(names) == 0 {
		t.Fatal("migrationFileNames() returned empty")
	}
	for _, name := range names {
		if filepath.Ext(name) != ".sql" {
			t.Fatalf("migrationFileNames() entry %q not .sql file", name)
		}
	}
}

func TestReadMigration(t *testing.T) {
	content, err := readMigration("0001_init.sql")
	if err != nil {
		t.Fatalf("readMigration('0001_init.sql') error = %v", err)
	}
	if len(content) == 0 {
		t.Fatal("readMigration() returned empty content")
	}
}

func TestReadMigrationNotFound(t *testing.T) {
	_, err := readMigration("nonexistent.sql")
	if err == nil {
		t.Fatal("readMigration('nonexistent.sql') error = nil, want error")
	}
}

func TestEnsureMigrationLedgerAndLoadAppliedMigrations(t *testing.T) {
	store := openTestDB(t)
	db := store.SQLDB()

	tx, err := db.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatalf("BeginTx() error = %v", err)
	}
	defer tx.Rollback()

	if err := ensureMigrationLedger(context.Background(), tx); err != nil {
		t.Fatalf("ensureMigrationLedger() error = %v", err)
	}
	if err := ensureMigrationLedger(context.Background(), tx); err != nil {
		t.Fatalf("ensureMigrationLedger() second call error = %v", err)
	}
	if _, err := tx.ExecContext(context.Background(), "INSERT INTO schema_migrations (version) VALUES (?)", "9999_test.sql"); err != nil {
		t.Fatalf("insert schema_migrations error = %v", err)
	}

	applied, err := loadAppliedMigrations(context.Background(), tx)
	if err != nil {
		t.Fatalf("loadAppliedMigrations() error = %v", err)
	}
	if !applied["9999_test.sql"] {
		t.Fatalf("loadAppliedMigrations() = %#v, want 9999_test.sql=true", applied)
	}
}

func TestBackfillLegacySchemaIfNeededRecordsCompleteLegacySchema(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "legacy.db")
	rawDB, err := sql.Open("sqlite", "file:"+filepath.ToSlash(dbPath))
	if err != nil {
		t.Fatalf("sql.Open() error = %v", err)
	}
	defer rawDB.Close()

	for _, ddl := range []string{
		"CREATE TABLE hosts (id INTEGER PRIMARY KEY, host_id TEXT NOT NULL, base_url TEXT NOT NULL, host_version TEXT NOT NULL, capability_probe_json TEXT NOT NULL DEFAULT '', created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)",
		"CREATE TABLE packs (id INTEGER PRIMARY KEY, pack_id TEXT NOT NULL, version TEXT NOT NULL, checksum TEXT NOT NULL, manifest_json TEXT NOT NULL DEFAULT '{}', metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)",
		"CREATE TABLE providers (id INTEGER PRIMARY KEY, pack_id INTEGER NOT NULL, provider_id TEXT NOT NULL, display_name TEXT NOT NULL, base_url TEXT NOT NULL, platform TEXT NOT NULL, account_type TEXT NOT NULL DEFAULT 'apikey', smoke_test_model TEXT NOT NULL DEFAULT '', provider_manifest_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)",
	} {
		if _, err := rawDB.ExecContext(context.Background(), ddl); err != nil {
			t.Fatalf("Exec legacy ddl error = %v", err)
		}
	}

	tx, err := rawDB.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatalf("BeginTx() error = %v", err)
	}
	defer tx.Rollback()

	if err := ensureMigrationLedger(context.Background(), tx); err != nil {
		t.Fatalf("ensureMigrationLedger() error = %v", err)
	}
	applied, err := loadAppliedMigrations(context.Background(), tx)
	if err != nil {
		t.Fatalf("loadAppliedMigrations() error = %v", err)
	}
	if err := backfillLegacySchemaIfNeeded(context.Background(), tx, []string{"0001_init.sql"}, applied); err != nil {
		t.Fatalf("backfillLegacySchemaIfNeeded() error = %v", err)
	}
	if !applied["0001_init.sql"] {
		t.Fatalf("applied = %#v, want 0001_init.sql marked", applied)
	}
}

func TestBackfillLegacySchemaIfNeededRejectsPartialLegacySchema(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "legacy-partial.db")
	rawDB, err := sql.Open("sqlite", "file:"+filepath.ToSlash(dbPath))
	if err != nil {
		t.Fatalf("sql.Open() error = %v", err)
	}
	defer rawDB.Close()

	if _, err := rawDB.ExecContext(context.Background(), "CREATE TABLE hosts (id INTEGER PRIMARY KEY)"); err != nil {
		t.Fatalf("Exec partial legacy ddl error = %v", err)
	}

	tx, err := rawDB.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatalf("BeginTx() error = %v", err)
	}
	defer tx.Rollback()

	if err := ensureMigrationLedger(context.Background(), tx); err != nil {
		t.Fatalf("ensureMigrationLedger() error = %v", err)
	}
	applied, err := loadAppliedMigrations(context.Background(), tx)
	if err != nil {
		t.Fatalf("loadAppliedMigrations() error = %v", err)
	}
	err = backfillLegacySchemaIfNeeded(context.Background(), tx, []string{"0001_init.sql"}, applied)
	if err == nil || err.Error() != "legacy sqlite schema is partially applied without schema_migrations" {
		t.Fatalf("backfillLegacySchemaIfNeeded() error = %v, want partial legacy schema error", err)
	}
}

func TestRollbackMigrationReturnsOriginalErrorWhenRollbackSucceeds(t *testing.T) {
	store := openTestDB(t)
	tx, err := store.SQLDB().BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatalf("BeginTx() error = %v", err)
	}

	baseErr := errors.New("apply failed")
	if err := rollbackMigration(tx, baseErr); !errors.Is(err, baseErr) {
		t.Fatalf("rollbackMigration() error = %v, want original error", err)
	}
}