package repository import ( "context" "encoding/json" "regexp" "strings" "testing" "time" "github.com/DATA-DOG/go-sqlmock" dbent "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/stretchr/testify/require" "entgo.io/ent/dialect" entsql "entgo.io/ent/dialect/sql" ) func newOllamaCloudUsageRepositoryTestClient(t *testing.T) (*dbent.Client, sqlmock.Sqlmock) { t.Helper() db, mock, err := sqlmock.New() require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) t.Cleanup(func() { _ = client.Close() }) return client, mock } func ollamaCloudUsageRepositoryAccount() *service.Account { return &service.Account{ ID: 17, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, Credentials: map[string]any{"api_key": "key", "base_url": "https://ollama.com"}, Extra: map[string]any{ service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=secret", service.OllamaCloudUsageAutoRefreshExtraKey: true, }, } } func TestUpdateOllamaCloudUsageSnapshotRowsAffectedZeroIsIdentityConflict(t *testing.T) { client, mock := newOllamaCloudUsageRepositoryTestClient(t) mock.ExpectBegin() expectOllamaCloudUsageGroupLock(mock, ollamaCloudUsageRepositoryAccount(), true, `"cipher:wos-session=secret"`, `true`, `null`) mock.ExpectExec(`(?s)`+regexp.QuoteMeta("UPDATE accounts")). WithArgs(sqlmock.AnyArg(), "key", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 0)) mock.ExpectRollback() repo := newAccountRepositoryWithSQL(client, nil, nil) err := repo.UpdateOllamaCloudUsageSnapshot(context.Background(), ollamaCloudUsageRepositoryAccount(), &service.OllamaCloudUsageSnapshot{ Status: service.OllamaCloudUsageStatusOK, LastAttemptAt: time.Now(), NextRefreshAt: time.Now().Add(time.Hour), }) require.ErrorIs(t, err, service.ErrOllamaCloudUsageIdentityChanged) require.NoError(t, mock.ExpectationsWereMet()) } func expectOllamaCloudUsageGroupLock( mock sqlmock.Sqlmock, account *service.Account, anchorMatches bool, sessionJSON, autoJSON, snapshotJSON string, ) { apiKey, _ := account.Credentials["api_key"].(string) credentials, _ := json.Marshal(normalizeJSONMap(account.Credentials)) var proxyID any if account.ProxyID != nil { proxyID = *account.ProxyID } mock.ExpectQuery(`(?s)`+regexp.QuoteMeta("SELECT")+`.*`+regexp.QuoteMeta("FOR NO KEY UPDATE")). WithArgs(apiKey, account.ID, account.Platform, account.Type, string(credentials), proxyID). WillReturnRows(sqlmock.NewRows([]string{"id", "anchor_matches", "session", "auto_refresh", "snapshot"}). AddRow(account.ID, anchorMatches, sessionJSON, autoJSON, snapshotJSON)) } func TestOllamaCloudUsageManagedWriteRejectsChangedProxyIdentity(t *testing.T) { client, mock := newOllamaCloudUsageRepositoryTestClient(t) mock.ExpectBegin() mock.ExpectQuery(`(?s)` + regexp.QuoteMeta("SELECT protocol, host, port") + `.*` + regexp.QuoteMeta("FOR SHARE")). WithArgs(int64(9)). WillReturnRows(sqlmock.NewRows([]string{"protocol", "host", "port", "username", "password", "status"}). AddRow("http", "new.example", 3128, "user", "pass", service.StatusActive)) mock.ExpectRollback() account := ollamaCloudUsageRepositoryAccount() proxyID := int64(9) account.ProxyID = &proxyID account.Proxy = &service.Proxy{ ID: proxyID, Protocol: "http", Host: "old.example", Port: 3128, Username: "user", Password: "pass", Status: service.StatusActive, } repo := newAccountRepositoryWithSQL(client, nil, nil) err := repo.SaveOllamaCloudUsageSession(context.Background(), account, "cipher:wos-session=replacement", true) require.ErrorIs(t, err, service.ErrOllamaCloudUsageIdentityChanged) require.NoError(t, mock.ExpectationsWereMet()) } func TestSaveAndDeleteOllamaCloudUsageSessionKeepCiphertextOutOfSQL(t *testing.T) { var capturedSQL []string matcher := sqlmock.QueryMatcherFunc(func(expectedSQL, actualSQL string) error { capturedSQL = append(capturedSQL, actualSQL) return sqlmock.QueryMatcherRegexp.Match(expectedSQL, actualSQL) }) db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(matcher)) require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) t.Cleanup(func() { _ = client.Close() }) repo := newAccountRepositoryWithSQL(client, db, nil) account := ollamaCloudUsageRepositoryAccount() const replacement = "cipher:wos-session=browser-cookie-secret" mock.ExpectBegin() expectOllamaCloudUsageGroupLock(mock, account, true, `"cipher:wos-session=secret"`, `true`, `null`) mock.ExpectExec(`(?s)UPDATE accounts.*ollama_cloud_usage_session.*ollama_cloud_usage_auto_refresh.*ollama_cloud_usage_snapshot`). WithArgs(`{"ollama_cloud_usage_auto_refresh":true,"ollama_cloud_usage_session":"cipher:wos-session=browser-cookie-secret"}`, "key", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() require.NoError(t, repo.SaveOllamaCloudUsageSession(context.Background(), account, replacement, true)) account.Extra[service.OllamaCloudUsageSessionExtraKey] = replacement mock.ExpectBegin() expectOllamaCloudUsageGroupLock(mock, account, true, `"cipher:wos-session=browser-cookie-secret"`, `true`, `null`) mock.ExpectExec(`(?s)UPDATE accounts.*ollama_cloud_usage_session.*ollama_cloud_usage_auto_refresh.*ollama_cloud_usage_snapshot`). WithArgs(`{}`, "key", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() require.NoError(t, repo.DeleteOllamaCloudUsageSession(context.Background(), account)) require.NotEmpty(t, capturedSQL) for _, query := range capturedSQL { require.NotContains(t, query, "browser-cookie-secret") } require.NoError(t, mock.ExpectationsWereMet()) } func TestOllamaCloudBaseURLSQLRegexMatchesServiceSemantics(t *testing.T) { for _, baseURL := range []string{ "https://ollama.com", "HTTPS://WWW.OLLAMA.COM:443/v1", "https://ollama.com/V1", "https://ollama.com/v1/", "https://ollama.com.evil.test/v1", } { t.Run(baseURL, func(t *testing.T) { matched, err := regexp.MatchString(ollamaCloudBaseURLRegexSQL, baseURL) require.NoError(t, err) account := ollamaCloudUsageRepositoryAccount() account.Credentials["base_url"] = baseURL require.Equal(t, service.IsOllamaCloudUsageAccount(account), matched) }) } } func TestListOllamaCloudUsageGroupAccountsUsesOneStrictBatchQuery(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) var capturedSQL string mock.ExpectQuery("SELECT id"). WithArgs(sqlmock.AnyArg()). WillReturnRows(sqlmock.NewRows([]string{"id"})) repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil) first := ollamaCloudUsageRepositoryAccount() second := ollamaCloudUsageRepositoryAccount() second.ID = 18 second.Platform = service.PlatformAnthropic second.Credentials = map[string]any{"api_key": "key", "base_url": "https://www.ollama.com:443/v1"} accounts, err := repo.ListOllamaCloudUsageGroupAccounts(context.Background(), []*service.Account{first, second}) require.NoError(t, err) require.Empty(t, accounts) query := normalizeSQLWhitespace(capturedSQL) require.Contains(t, query, "credentials ->> 'api_key' = ANY($1)") require.Contains(t, query, "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek', 'minimax')") require.Contains(t, query, "jsonb_typeof(credentials -> 'api_key') = 'string'") require.Contains(t, query, ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")) require.NotContains(t, query, "~*") require.NoError(t, mock.ExpectationsWereMet()) } func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) debounce := time.Minute maxWait := time.Hour var capturedSQL string mock.ExpectQuery("WITH eligible AS"). WithArgs(now.UTC(), debounce.Seconds(), maxWait.Seconds(), 20, service.OllamaCloudUsageMinFetchInterval.Seconds()). WillReturnRows(sqlmock.NewRows([]string{"id", "group_last_used_at"})) repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil) accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, debounce, maxWait, 20) require.NoError(t, err) require.Empty(t, accounts) normalized := normalizeSQLWhitespace(capturedSQL) for _, clause := range []string{ "deleted_at IS NULL", "status = 'active'", "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek', 'minimax')", "type = 'apikey'", ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'"), "jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'", `extra @> '{"ollama_cloud_usage_auto_refresh": true}'::jsonb`, "MAX(last_used_at) AS group_last_used_at", "PARTITION BY api_key", "WHERE group_rank = 1", "LIMIT $4", "make_interval(secs => $2::double precision)", "make_interval(secs => $3::double precision)", // Minimum interval floor between successful fetches. "make_interval(secs => $5::double precision)", // jsonpath .datetime() only accepts the ISO-8601 "Z" designator from // PostgreSQL 17 on, and this service writes UTC timestamps. Without this // rewrite every parsed_* column is NULL on 14-16 and the due filter // collapses into its fail-open branch. `regexp_replace( regexp_replace( fetched_at, '(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$', '\1\2' ), 'Z$', '+00:00' )`, "group_last_used_at > parsed_fetched_at::timestamptz", "group_last_used_at > parsed_last_attempt_at::timestamptz", "$1 >= activity_due_at", "COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)", "ORDER BY due_class, due_at NULLS FIRST, id", } { require.Contains(t, normalized, clause) } require.NotContains(t, normalized, "~*") require.NoError(t, mock.ExpectationsWereMet()) } func TestBulkUpdateOllamaIdentityCleanupIsValueConditional(t *testing.T) { exec := &recordingSQLExecutor{result: rowsAffectedResult(1)} repo := newAccountRepositoryWithSQL(nil, exec, nil) _, err := repo.BulkUpdate(context.Background(), []int64{17}, service.AccountBulkUpdate{ Credentials: map[string]any{"base_url": "https://www.ollama.com:443/v1"}, }) require.NoError(t, err) require.NotEmpty(t, exec.execQueries) query := normalizeSQLWhitespace(exec.execQueries[0]) require.Contains(t, query, "NOT ("+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")) require.Contains(t, query, ollamaCloudBaseURLMatchesSQL("$1::jsonb ->> 'base_url'")) require.NotContains(t, query, "~*") require.Contains(t, query, "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek', 'minimax') AND type = 'apikey'") require.Contains(t, query, "- 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'") payload, ok := exec.execArgs[0][0].([]byte) require.True(t, ok) require.NotContains(t, string(payload), service.OllamaCloudUsageSnapshotExtraKey) } func TestUpdateCredentialsIdentityChangeClearsAllOllamaManagedExtra(t *testing.T) { client, mock := newOllamaCloudUsageRepositoryTestClient(t) mock.ExpectBegin() mock.ExpectExec(`(?s)UPDATE accounts.*credentials -> 'api_key' IS DISTINCT FROM.*ollama_cloud_usage_session.*ollama_cloud_usage_auto_refresh.*ollama_cloud_usage_snapshot`). WithArgs(`{"api_key":"new-key","base_url":"https://ollama.com"}`, int64(17)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")). WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectCommit() repo := newAccountRepositoryWithSQL(client, nil, nil) err := repo.UpdateCredentials(context.Background(), 17, map[string]any{ "api_key": "new-key", "base_url": "https://ollama.com", }) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } func TestDisableOllamaCloudUsageAutoRefreshUsesGroupIdentityCAS(t *testing.T) { client, mock := newOllamaCloudUsageRepositoryTestClient(t) account := ollamaCloudUsageRepositoryAccount() mock.ExpectBegin() expectOllamaCloudUsageGroupLock(mock, account, true, `"cipher:wos-session=secret"`, `true`, `null`) mock.ExpectExec(`(?s)UPDATE accounts.*ollama_cloud_usage_auto_refresh`). WithArgs(`{"ollama_cloud_usage_auto_refresh":false,"ollama_cloud_usage_session":"cipher:wos-session=secret"}`, "key", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() repo := newAccountRepositoryWithSQL(client, nil, nil) err := repo.DisableOllamaCloudUsageAutoRefresh(context.Background(), account) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } // Ollama 清理分支必须带顶层 credentials DISTINCT 守卫:没有它,非 Ollama 的 // openai/anthropic apikey 账号在凭证未变化的持久化上也会误清探测快照。 func TestUpdateCredentialsCleanupBranchRequiresChangedCredentials(t *testing.T) { client, mock := newOllamaCloudUsageRepositoryTestClient(t) mock.ExpectBegin() mock.ExpectExec(`(?s)UPDATE accounts.*CASE.*AND credentials IS DISTINCT FROM \$1::jsonb\s+AND \(\s+credentials -> 'api_key' IS DISTINCT FROM`). WithArgs(`{"api_key":"same-key","base_url":"https://relay.example.com/v1"}`, int64(17)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")). WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectCommit() repo := newAccountRepositoryWithSQL(client, nil, nil) err := repo.UpdateCredentials(context.Background(), 17, map[string]any{ "api_key": "same-key", "base_url": "https://relay.example.com/v1", }) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } // SQL 平台白名单常量与 service 判定必须互为镜像:对每个已知平台,常量里的 // 成员关系都要与 IsOllamaCloudUsageAccount(apikey + 官方 ollama.com)一致, // 防止两侧平台列表各自漂移。 func TestOllamaCloudUsagePlatformWhitelistMatchesServicePredicate(t *testing.T) { matches := regexp.MustCompile(`'([^']+)'`).FindAllStringSubmatch(ollamaCloudUsagePlatformsSQL, -1) sqlPlatforms := make(map[string]struct{}, len(matches)) for _, match := range matches { sqlPlatforms[match[1]] = struct{}{} } require.Len(t, sqlPlatforms, 6) for _, platform := range []string{ service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformGemini, service.PlatformGrok, service.PlatformAntigravity, service.PlatformComposite, "kiro", } { account := ollamaCloudUsageRepositoryAccount() account.Platform = platform _, inSQL := sqlPlatforms[platform] require.Equal(t, inSQL, service.IsOllamaCloudUsageAccount(account), platform) } } // 语义等价性:平台放开后,普通(非 ollama)kimi apikey 账号改凭证会从通用 // apikey 分支落到 Ollama 分支——多减三个它本来就不存在的 ollama 键。分支必须 // 仍然清掉 probe 快照,且不得减其它任何 extra 键。 func TestUpdateCredentialsPlainCNAPIKeyAccountCleanupStaysSemanticallyEquivalent(t *testing.T) { var capturedSQL string matcher := sqlmock.QueryMatcherFunc(func(expectedSQL, actualSQL string) error { if strings.Contains(actualSQL, "UPDATE accounts") { capturedSQL = actualSQL } return sqlmock.QueryMatcherRegexp.Match(expectedSQL, actualSQL) }) db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(matcher)) require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) t.Cleanup(func() { _ = client.Close() }) mock.ExpectBegin() mock.ExpectExec(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_session'.*- 'ollama_cloud_usage_auto_refresh'.*- 'ollama_cloud_usage_snapshot'`). WithArgs(`{"api_key":"rotated-key","base_url":"https://api.moonshot.cn/v1"}`, int64(17)). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")). WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectCommit() repo := newAccountRepositoryWithSQL(client, db, nil) err = repo.UpdateCredentials(context.Background(), 17, map[string]any{ "api_key": "rotated-key", "base_url": "https://api.moonshot.cn/v1", }) require.NoError(t, err) query := normalizeSQLWhitespace(capturedSQL) require.Contains(t, query, "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek', 'minimax') AND type = 'apikey' AND credentials IS DISTINCT FROM $1::jsonb") require.Contains(t, query, "THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe' - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'") require.NotContains(t, query, "- 'upstream_billing_probe_enabled'") require.NotContains(t, query, "- 'upstream_billing_rate_sync_enabled'") require.NoError(t, mock.ExpectationsWereMet()) }