sub2api-cn-relay-manager/internal/provision/runtime_import_service_test.go

package provision

import (
	"context"
	"database/sql"
	"fmt"
	"path/filepath"
	"testing"

	_ "modernc.org/sqlite"
	"sub2api-cn-relay-manager/internal/host/sub2api"
	"sub2api-cn-relay-manager/internal/pack"
	"sub2api-cn-relay-manager/internal/store/sqlite"
)

func TestRuntimeImportServicePersistsOperationalState(t *testing.T) {
	store := openProvisionTestStore(t)
	defer closeProvisionTestStore(t, store)

	seedProvisionHost(t, store, "host-1", "https://sub2api.example.com")

	host := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_1"}, {ID: "account_2"}},
		testResults: map[string]sub2api.ProbeResult{
			"account_1": {OK: true, Status: "passed"},
			"account_2": {OK: true, Status: "passed"},
		},
		models: map[string][]sub2api.AccountModel{
			"account_1": {{ID: "deepseek-chat"}},
			"account_2": {{ID: "deepseek-chat"}},
		},
		gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}},
	}

	svc := NewRuntimeImportService(store, host)
	result, err := svc.Import(context.Background(), RuntimeImportRequest{
		HostID:      "host-1",
		HostBaseURL: "https://sub2api.example.com",
		Pack: pack.LoadedPack{
			Manifest: pack.Manifest{PackID: "openai-cn-pack", Version: "1.0.0", TargetHost: "sub2api", MinHostVersion: "0.1.126", MaxHostVersion: "0.2.x"},
			Checksum: "checksum-1",
		},
		Provider: sampleProviderManifest(),
		Mode:     ImportModePartial,
		Keys:     []string{" key-1 ", "key-2", "key-1"},
		Access: AccessRequest{
			Mode:        AccessModeSelfService,
			ProbeAPIKey: "user-key",
		},
	})
	if err != nil {
		t.Fatalf("RuntimeImportService.Import() error = %v", err)
	}
	if result.BatchID <= 0 {
		t.Fatalf("BatchID = %d, want positive id", result.BatchID)
	}
	if result.Report.BatchStatus != BatchStatusSucceeded {
		t.Fatalf("BatchStatus = %q, want %q", result.Report.BatchStatus, BatchStatusSucceeded)
	}

	if got := queryCount(t, store.SQLDB(), "hosts"); got != 1 {
		t.Fatalf("hosts row count = %d, want 1", got)
	}
	if got := queryCount(t, store.SQLDB(), "packs"); got != 1 {
		t.Fatalf("packs row count = %d, want 1", got)
	}
	if got := queryCount(t, store.SQLDB(), "providers"); got != 1 {
		t.Fatalf("providers row count = %d, want 1", got)
	}
	if got := queryCount(t, store.SQLDB(), "import_batches"); got != 1 {
		t.Fatalf("import_batches row count = %d, want 1", got)
	}
	if got := queryCount(t, store.SQLDB(), "import_batch_items"); got != 2 {
		t.Fatalf("import_batch_items row count = %d, want 2", got)
	}
	if got := queryCount(t, store.SQLDB(), "managed_resources"); got != 4 {
		t.Fatalf("managed_resources row count = %d, want 4", got)
	}
	if got := queryCount(t, store.SQLDB(), "probe_results"); got != 2 {
		t.Fatalf("probe_results row count = %d, want 2", got)
	}
	if got := queryCount(t, store.SQLDB(), "access_closure_records"); got != 1 {
		t.Fatalf("access_closure_records row count = %d, want 1", got)
	}

	var batchStatus string
	var accessStatus string
	if err := store.SQLDB().QueryRowContext(context.Background(), "SELECT batch_status, access_status FROM import_batches WHERE id = ?", result.BatchID).Scan(&batchStatus, &accessStatus); err != nil {
		t.Fatalf("query import batch state: %v", err)
	}
	if batchStatus != BatchStatusSucceeded {
		t.Fatalf("persisted batch_status = %q, want %q", batchStatus, BatchStatusSucceeded)
	}
	if accessStatus != AccessStatusSelfServiceReady {
		t.Fatalf("persisted access_status = %q, want %q", accessStatus, AccessStatusSelfServiceReady)
	}

	var fingerprint string
	var accountStatus string
	if err := store.SQLDB().QueryRowContext(context.Background(), "SELECT key_fingerprint, account_status FROM import_batch_items ORDER BY id LIMIT 1").Scan(&fingerprint, &accountStatus); err != nil {
		t.Fatalf("query import batch item: %v", err)
	}
	if fingerprint == "key-1" || fingerprint == "key-2" || len(fingerprint) < 10 {
		t.Fatalf("key_fingerprint = %q, want hashed fingerprint instead of raw key", fingerprint)
	}
	if accountStatus != "passed" {
		t.Fatalf("account_status = %q, want passed", accountStatus)
	}
}

func TestRuntimeImportServicePersistsFailedBatchAfterStrictRollback(t *testing.T) {
	store := openProvisionTestStore(t)
	defer closeProvisionTestStore(t, store)

	seedProvisionHost(t, store, "host-1", "https://sub2api.example.com")

	host := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_1"}, {ID: "account_2"}},
		testResults: map[string]sub2api.ProbeResult{
			"account_1": {OK: true, Status: "passed"},
			"account_2": {OK: false, Status: "failed", Message: "bad key"},
		},
		models: map[string][]sub2api.AccountModel{
			"account_1": {{ID: "deepseek-chat"}},
			"account_2": {{ID: "deepseek-chat"}},
		},
	}

	svc := NewRuntimeImportService(store, host)
	result, err := svc.Import(context.Background(), RuntimeImportRequest{
		HostID:      "host-1",
		HostBaseURL: "https://sub2api.example.com",
		Pack: pack.LoadedPack{
			Manifest: pack.Manifest{PackID: "openai-cn-pack", Version: "1.0.0", TargetHost: "sub2api", MinHostVersion: "0.1.126", MaxHostVersion: "0.2.x"},
			Checksum: "checksum-1",
		},
		Provider: sampleProviderManifest(),
		Mode:     ImportModeStrict,
		Keys:     []string{"key-1", "key-2"},
		Access: AccessRequest{
			Mode:        AccessModeSelfService,
			ProbeAPIKey: "user-key",
		},
	})
	if err == nil {
		t.Fatal("RuntimeImportService.Import() error = nil, want strict failure")
	}
	if result.BatchID <= 0 {
		t.Fatalf("BatchID = %d, want positive id", result.BatchID)
	}

	var batchStatus string
	var accessStatus string
	if err := store.SQLDB().QueryRowContext(context.Background(), "SELECT batch_status, access_status FROM import_batches WHERE id = ?", result.BatchID).Scan(&batchStatus, &accessStatus); err != nil {
		t.Fatalf("query failed import batch state: %v", err)
	}
	if batchStatus != BatchStatusFailed {
		t.Fatalf("persisted batch_status = %q, want %q", batchStatus, BatchStatusFailed)
	}
	if accessStatus != AccessStatusBroken {
		t.Fatalf("persisted access_status = %q, want %q", accessStatus, AccessStatusBroken)
	}
	if got := queryCount(t, store.SQLDB(), "managed_resources"); got != 0 {
		t.Fatalf("managed_resources row count = %d, want 0 after strict rollback", got)
	}
	if got := queryCount(t, store.SQLDB(), "probe_results"); got != 2 {
		t.Fatalf("probe_results row count = %d, want 2", got)
	}
	if got := queryCount(t, store.SQLDB(), "access_closure_records"); got != 1 {
		t.Fatalf("access_closure_records row count = %d, want 1", got)
	}
}

func TestRuntimeImportServicePersistsPartialManagedResourcesOnAccessFailure(t *testing.T) {
	store := openProvisionTestStore(t)
	defer closeProvisionTestStore(t, store)

	seedProvisionHost(t, store, "host-1", "https://sub2api.example.com")

	host := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_1"}},
		testResults: map[string]sub2api.ProbeResult{
			"account_1": {OK: true, Status: "passed"},
		},
		models: map[string][]sub2api.AccountModel{
			"account_1": {{ID: "deepseek-chat"}},
		},
		assignErr: fmt.Errorf("group is not a subscription type"),
	}

	svc := NewRuntimeImportService(store, host)
	result, err := svc.Import(context.Background(), RuntimeImportRequest{
		HostID:      "host-1",
		HostBaseURL: "https://sub2api.example.com",
		Pack: pack.LoadedPack{
			Manifest: pack.Manifest{PackID: "openai-cn-pack", Version: "1.0.0", TargetHost: "sub2api", MinHostVersion: "0.1.126", MaxHostVersion: "0.2.x"},
			Checksum: "checksum-1",
		},
		Provider: sampleProviderManifest(),
		Mode:     ImportModePartial,
		Keys:     []string{"key-1"},
		Access: AccessRequest{
			Mode:          AccessModeSubscription,
			ProbeAPIKey:   "user-key",
			Subscriptions: []SubscriptionTarget{{UserID: "1", DurationDays: 30}},
		},
	})
	if err == nil {
		t.Fatal("RuntimeImportService.Import() error = nil, want partial failure")
	}
	if result.BatchID <= 0 {
		t.Fatalf("BatchID = %d, want positive id", result.BatchID)
	}
	if got := queryCount(t, store.SQLDB(), "managed_resources"); got != 4 {
		t.Fatalf("managed_resources row count = %d, want 4 persisted resources on partial failure", got)
	}
	var batchStatus string
	if err := store.SQLDB().QueryRowContext(context.Background(), "SELECT batch_status FROM import_batches WHERE id = ?", result.BatchID).Scan(&batchStatus); err != nil {
		t.Fatalf("query import batch status: %v", err)
	}
	if batchStatus != BatchStatusPartial {
		t.Fatalf("persisted batch_status = %q, want %q", batchStatus, BatchStatusPartial)
	}
}

func TestRuntimeImportServiceRepeatedImportReusesManagedResources(t *testing.T) {
	store := openProvisionTestStore(t)
	defer closeProvisionTestStore(t, store)

	seedProvisionHost(t, store, "host-1", "https://sub2api.example.com")

	host := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_1", Name: "key-1"}},
		testResults: map[string]sub2api.ProbeResult{
			"account_1": {OK: true, Status: "passed"},
		},
		models: map[string][]sub2api.AccountModel{
			"account_1": {{ID: "deepseek-chat"}},
		},
		gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}},
	}

	svc := NewRuntimeImportService(store, host)
	request := RuntimeImportRequest{
		HostID:      "host-1",
		HostBaseURL: "https://sub2api.example.com",
		Pack: pack.LoadedPack{
			Manifest: pack.Manifest{PackID: "openai-cn-pack", Version: "1.0.0", TargetHost: "sub2api", MinHostVersion: "0.1.126", MaxHostVersion: "0.2.x"},
			Checksum: "checksum-1",
		},
		Provider: sampleProviderManifest(),
		Mode:     ImportModePartial,
		Keys:     []string{"key-1"},
		Access: AccessRequest{
			Mode:        AccessModeSelfService,
			ProbeAPIKey: "user-key",
		},
	}

	first, err := svc.Import(context.Background(), request)
	if err != nil {
		t.Fatalf("first Import() error = %v", err)
	}
	second, err := svc.Import(context.Background(), request)
	if err != nil {
		t.Fatalf("second Import() error = %v", err)
	}
	if second.BatchID <= first.BatchID {
		t.Fatalf("second BatchID = %d, want > first BatchID %d", second.BatchID, first.BatchID)
	}
	if got := queryCount(t, store.SQLDB(), "managed_resources"); got != 3 {
		t.Fatalf("managed_resources row count = %d, want 3 after reused import", got)
	}
	if got := queryCount(t, store.SQLDB(), "import_batches"); got != 2 {
		t.Fatalf("import_batches row count = %d, want 2", got)
	}
}

func openProvisionTestStore(t *testing.T) *sqlite.DB {
	t.Helper()

	dbPath := filepath.Join(t.TempDir(), "state.db")
	dsn := fmt.Sprintf("file:%s?_busy_timeout=5000&_pragma=foreign_keys(0)", filepath.ToSlash(dbPath))
	store, err := sqlite.Open(context.Background(), dsn)
	if err != nil {
		t.Fatalf("sqlite.Open() error = %v", err)
	}
	return store
}

func closeProvisionTestStore(t *testing.T, store *sqlite.DB) {
	t.Helper()
	if err := store.Close(); err != nil {
		t.Fatalf("store.Close() error = %v", err)
	}
}

func seedProvisionHost(t *testing.T, store *sqlite.DB, hostID, baseURL string) int64 {
	t.Helper()
	id, err := store.Hosts().Create(context.Background(), sqlite.Host{
		HostID:      hostID,
		BaseURL:     baseURL,
		HostVersion: "0.1.126",
		AuthType:    "apikey",
		AuthToken:   "test-host-token",
	})
	if err != nil {
		t.Fatalf("Hosts().Create() error = %v", err)
	}
	return id
}

func queryCount(t *testing.T, db *sql.DB, table string) int {
	t.Helper()
	var count int
	if err := db.QueryRowContext(context.Background(), "SELECT COUNT(*) FROM "+table).Scan(&count); err != nil {
		t.Fatalf("count rows for %s: %v", table, err)
	}
	return count
}