309 lines
14 KiB
Go
309 lines
14 KiB
Go
package localstore_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/administration"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/identity"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/localstore"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/templates"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage"
|
|
)
|
|
|
|
var fixed = time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
|
|
func TestStoreSatisfiesApplicationPersistenceSeams(t *testing.T) {
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
var _ administration.Store = store
|
|
var _ assets.Catalog = store
|
|
var _ billing.Store = store
|
|
var _ billing.PriceRuleSeeder = store
|
|
var _ billing.WalletPoster = store
|
|
var _ jobs.Store = store
|
|
var _ templates.Catalog = store
|
|
var _ usage.Repository = store
|
|
var _ identity.AuthorizationSnapshotLoader = store
|
|
var _ identity.CredentialAuthenticator = store
|
|
var _ identity.PasswordChanger = store
|
|
var _ interface {
|
|
GetOwnerByStoragePath(context.Context, string, string) (assets.Asset, bool, error)
|
|
FailCreation(context.Context, jobs.Job) error
|
|
WriteBilling(context.Context, string, json.RawMessage) error
|
|
WriteOutputAssetIDs(context.Context, string, []string) error
|
|
} = store
|
|
}
|
|
|
|
func TestAssetLifecycleIsOwnerScopedAndPublicAccessUsesTagsOrJobs(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New()
|
|
a := assets.Asset{ID: "asset-1", OwnerID: "owner-a", StoragePath: "uploads/a.png", Tags: []string{assets.ClientTag("client-a")}, Metadata: map[string]any{"x": "y"}, CreatedAt: fixed}
|
|
if _, err := store.Create(ctx, a); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, found, _ := store.GetOwner(ctx, "owner-b", a.ID); found {
|
|
t.Fatal("cross-owner asset leaked")
|
|
}
|
|
if got, found, _ := store.GetOwnerByStoragePath(ctx, a.OwnerID, a.StoragePath); !found || got.ID != a.ID {
|
|
t.Fatalf("storage lookup = %#v, %v", got, found)
|
|
}
|
|
if got, _ := store.ListPublic(ctx, a.OwnerID, "client-a", 10); len(got) != 1 {
|
|
t.Fatalf("public tagged assets = %d", len(got))
|
|
}
|
|
if _, found, _ := store.DeleteOwner(ctx, "owner-b", a.ID); found {
|
|
t.Fatal("cross-owner delete succeeded")
|
|
}
|
|
}
|
|
|
|
func TestJobClaimsRespectScheduleLeaseAndOrchestrationWrites(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
ready := fixed.Add(-time.Minute)
|
|
job := jobs.Job{ID: "job-1", OwnerID: "owner", ExternalClientID: "client", Status: jobs.StatusQueued, Capability: "image.generate", Provider: "mock", ReqKey: "mock", ScheduledAt: fixed.Add(-time.Minute), DispatchReadyAt: &ready, CreatedAt: fixed.Add(-time.Hour), UpdatedAt: fixed.Add(-time.Hour)}
|
|
if _, err := store.CreateJob(ctx, job); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
claimed, err := store.ClaimJobs(ctx, "worker-a", 1, 300)
|
|
if err != nil || len(claimed) != 1 || claimed[0].Status != jobs.StatusQueued || claimed[0].LockedBy != "worker-a" || claimed[0].Attempts != 0 {
|
|
t.Fatalf("claim = %#v, %v", claimed, err)
|
|
}
|
|
if again, _ := store.ClaimJobs(ctx, "worker-b", 1, 300); len(again) != 0 {
|
|
t.Fatal("active lease was reclaimed")
|
|
}
|
|
if err := store.WriteBilling(ctx, job.ID, json.RawMessage(`{"amountFen":12}`)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.WriteOutputAssetIDs(ctx, job.ID, []string{"asset-out"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
failed := job
|
|
failed.Status = jobs.StatusFailed
|
|
failed.Error = &jobs.JobError{Message: "charge failed"}
|
|
failed.Billing = json.RawMessage(`{"status":"failed"}`)
|
|
if err := store.FailCreation(ctx, failed); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, found, _ := store.FindJob(ctx, jobs.Scope{OwnerID: "owner"}, job.ID)
|
|
if !found || got.Status != jobs.StatusFailed || got.LockedBy != "" || got.OutputAssetIDs[0] != "asset-out" {
|
|
t.Fatalf("job = %#v", got)
|
|
}
|
|
}
|
|
|
|
func TestWalletPostingIsIdempotentAndNeverOverdrafts(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
org := administration.Organization{ID: "org", Name: "Org", Status: administration.StatusActive, CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateOrganization(ctx, org); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
credit := billing.WalletPostParams{LedgerID: "l1", OrganizationID: org.ID, Kind: "recharge", DeltaFen: 100, Currency: billing.CurrencyCNY, IdempotencyKey: "credit-1"}
|
|
first, err := store.PostWalletEntry(ctx, credit)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
repeated, err := store.PostWalletEntry(ctx, credit)
|
|
if err != nil || repeated.LedgerID != first.LedgerID {
|
|
t.Fatalf("repeat = %#v, %v", repeated, err)
|
|
}
|
|
changed := credit
|
|
changed.DeltaFen = 99
|
|
if _, err := store.PostWalletEntry(ctx, changed); err == nil {
|
|
t.Fatal("payload mismatch accepted")
|
|
}
|
|
if _, err := store.PostWalletEntry(ctx, billing.WalletPostParams{LedgerID: "l2", OrganizationID: org.ID, Kind: "charge", DeltaFen: -101, Currency: billing.CurrencyCNY, IdempotencyKey: "charge-1"}); err == nil {
|
|
t.Fatal("overdraft accepted")
|
|
}
|
|
wallet, _ := store.BillingWallet(ctx, org.ID)
|
|
if wallet.BalanceFen != 100 || wallet.TotalRechargedFen != 100 {
|
|
t.Fatalf("wallet = %#v", wallet)
|
|
}
|
|
}
|
|
|
|
func TestWalletPostingIdempotencyIgnoresDescriptionMetadataAndLedgerID(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
params := billing.WalletPostParams{
|
|
LedgerID: "ledger-original", OrganizationID: "org", AccountID: "account", JobID: "job",
|
|
Kind: "recharge", DeltaFen: 100, Currency: billing.CurrencyCNY, IdempotencyKey: "credit",
|
|
Description: "original", Metadata: map[string]any{"request": "original"},
|
|
}
|
|
first, err := store.PostWalletEntry(ctx, params)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
replay := params
|
|
replay.LedgerID = "ledger-replay"
|
|
replay.Description = "changed"
|
|
replay.Metadata = map[string]any{"request": "changed", "extra": true}
|
|
got, err := store.PostWalletEntry(ctx, replay)
|
|
if err != nil {
|
|
t.Fatalf("replay with non-conflicting fields: %v", err)
|
|
}
|
|
if got.LedgerID != first.LedgerID {
|
|
t.Fatalf("replay ledger = %q, want original %q", got.LedgerID, first.LedgerID)
|
|
}
|
|
ledger, err := store.BillingLedger(ctx, "org", "", 10)
|
|
if err != nil || len(ledger) != 1 {
|
|
t.Fatalf("ledger = %#v, %v", ledger, err)
|
|
}
|
|
}
|
|
|
|
func TestActivateCreationReplayRequiresSemanticallyEqualBilling(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
job := jobs.Job{ID: "activation", OwnerID: "owner", Status: jobs.StatusQueued, ScheduledAt: fixed, CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateJob(ctx, job); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.ActivateCreation(ctx, job.ID, json.RawMessage(`{"status":"charged","details":{"amountFen":12,"currency":"CNY"}}`)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.ActivateCreation(ctx, job.ID, json.RawMessage(`{ "details": { "currency": "CNY", "amountFen": 12 }, "status": "charged" }`)); err != nil {
|
|
t.Fatalf("semantic replay: %v", err)
|
|
}
|
|
if err := store.ActivateCreation(ctx, job.ID, json.RawMessage(`{"status":"charged","details":{"amountFen":13,"currency":"CNY"}}`)); !errors.Is(err, jobs.ErrTransitionConflict) {
|
|
t.Fatalf("different replay error = %v", err)
|
|
}
|
|
got, found, err := store.FindJob(ctx, jobs.Scope{OwnerID: job.OwnerID}, job.ID)
|
|
if err != nil || !found || string(got.Billing) != `{"status":"charged","details":{"amountFen":12,"currency":"CNY"}}` {
|
|
t.Fatalf("job after conflict = %#v, %v, %v", got, found, err)
|
|
}
|
|
}
|
|
|
|
func TestDeleteAccountArchivesUsageEvents(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
account := administration.Account{ID: "delete-me", Phone: "13900139000", OrganizationID: localstore.DemoOrganizationID, Status: administration.StatusActive, CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateAccount(ctx, account); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
event := usage.Event{ID: "usage-delete", JobID: "job-delete", OwnerID: account.ID, OrganizationID: account.OrganizationID, CreatedAt: fixed.Format(time.RFC3339)}
|
|
if _, inserted, err := store.Insert(event); err != nil || !inserted {
|
|
t.Fatalf("insert usage = %v, %v", inserted, err)
|
|
}
|
|
if err := store.DeleteAccount(ctx, account.ID, "archive-owner"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
archived, err := store.List(usage.Filters{OwnerID: "archive-owner"})
|
|
if err != nil || len(archived) != 1 || archived[0].ID != event.ID {
|
|
t.Fatalf("archived usage = %#v, %v", archived, err)
|
|
}
|
|
original, err := store.List(usage.Filters{OwnerID: account.ID})
|
|
if err != nil || len(original) != 0 {
|
|
t.Fatalf("original usage = %#v, %v", original, err)
|
|
}
|
|
}
|
|
|
|
func TestGenerationLifecycleRejectsStaleCASAndFencedWriters(t *testing.T) {
|
|
ctx := context.Background()
|
|
now := fixed
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return now }))
|
|
ready := fixed
|
|
job := jobs.Job{ID: "fenced", OwnerID: "owner", Status: jobs.StatusQueued, ScheduledAt: fixed, DispatchReadyAt: &ready, CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateJob(ctx, job); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
first, err := store.ClaimJobs(ctx, "worker-old", 1, 60)
|
|
if err != nil || len(first) != 1 {
|
|
t.Fatalf("first claim = %#v, %v", first, err)
|
|
}
|
|
now = fixed.Add(2 * time.Minute)
|
|
second, err := store.ClaimJobs(ctx, "worker-new", 1, 60)
|
|
if err != nil || len(second) != 1 || second[0].LockedBy != "worker-new" {
|
|
t.Fatalf("second claim = %#v, %v", second, err)
|
|
}
|
|
oldWorker := "worker-old"
|
|
running := jobs.StatusRunning
|
|
if _, err := store.UpdateJob(ctx, job.ID, jobs.Patch{Status: &running, ExpectedStatuses: []jobs.Status{jobs.StatusQueued}, ExpectedLockedBy: &oldWorker}); !errors.Is(err, jobs.ErrTransitionConflict) {
|
|
t.Fatalf("stale CAS error = %v", err)
|
|
}
|
|
if err := store.WriteBillingFenced(ctx, job.ID, json.RawMessage(`{"status":"old"}`), jobs.StatusQueued, oldWorker); !errors.Is(err, jobs.ErrTransitionConflict) {
|
|
t.Fatalf("stale billing writer error = %v", err)
|
|
}
|
|
if err := store.WriteOutputAssetIDsFenced(ctx, job.ID, []string{"old-output"}, jobs.StatusQueued, oldWorker); !errors.Is(err, jobs.ErrTransitionConflict) {
|
|
t.Fatalf("stale output writer error = %v", err)
|
|
}
|
|
newWorker := "worker-new"
|
|
updated, err := store.UpdateJob(ctx, job.ID, jobs.Patch{Status: &running, ExpectedStatuses: []jobs.Status{jobs.StatusQueued}, ExpectedLockedBy: &newWorker})
|
|
if err != nil || updated.Status != jobs.StatusRunning {
|
|
t.Fatalf("current CAS update = %#v, %v", updated, err)
|
|
}
|
|
if err := store.WriteBillingFenced(ctx, job.ID, json.RawMessage(`{"status":"current"}`), jobs.StatusRunning, newWorker); err != nil {
|
|
t.Fatalf("current billing writer: %v", err)
|
|
}
|
|
if err := store.WriteOutputAssetIDsFenced(ctx, job.ID, []string{"current-output"}, jobs.StatusRunning, newWorker); err != nil {
|
|
t.Fatalf("current output writer: %v", err)
|
|
}
|
|
got, found, err := store.FindJob(ctx, jobs.Scope{OwnerID: job.OwnerID}, job.ID)
|
|
if err != nil || !found || got.LockedBy != newWorker || got.Status != jobs.StatusRunning || string(got.Billing) != `{"status":"current"}` || len(got.OutputAssetIDs) != 1 || got.OutputAssetIDs[0] != "current-output" {
|
|
t.Fatalf("job after stale writes = %#v, %v, %v", got, found, err)
|
|
}
|
|
}
|
|
|
|
func TestAdministrationIdentityAndPasswordLifecycleShareState(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
seed, found, err := store.GetAccount(ctx, localstore.DemoAccountID)
|
|
if err != nil || !found || seed.Role != administration.RoleSuperAdmin || seed.OrganizationID != localstore.DemoOrganizationID {
|
|
t.Fatalf("demo = %#v, %v, %v", seed, found, err)
|
|
}
|
|
hashed, _ := administration.HashPassword("old-password")
|
|
account := administration.Account{ID: "user-1", Phone: "13800138000", DisplayName: "User", Role: administration.RoleUser, OrganizationID: localstore.DemoOrganizationID, Status: administration.StatusActive, PasswordHash: hashed.Hash, PasswordSalt: hashed.Salt, SessionVersion: 1, CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateAccount(ctx, account); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := store.AttemptPasswordLogin(ctx, account.Phone, "wrong", fixed); !errors.Is(err, identity.ErrPasswordLogin) {
|
|
t.Fatalf("wrong password error = %v", err)
|
|
}
|
|
login, err := store.AttemptPasswordLogin(ctx, account.Phone, "old-password", fixed)
|
|
if err != nil || login.Account.ID != account.ID {
|
|
t.Fatalf("login = %#v, %v", login, err)
|
|
}
|
|
snapshot, err := store.ChangeOwnPassword(ctx, account.ID, "old-password", "new-password", fixed.Add(time.Hour))
|
|
if err != nil || snapshot.Account.SessionVersion != 2 {
|
|
t.Fatalf("change = %#v, %v", snapshot, err)
|
|
}
|
|
if _, err := store.AttemptPasswordLogin(ctx, account.Phone, "old-password", fixed.Add(time.Hour)); !errors.Is(err, identity.ErrPasswordLogin) {
|
|
t.Fatal("old password still works")
|
|
}
|
|
if _, err := store.AttemptPasswordLogin(ctx, account.Phone, "new-password", fixed.Add(time.Hour)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestTemplatesUsageAndPriceSeedingAreIdempotentAndScoped(t *testing.T) {
|
|
ctx := context.Background()
|
|
store := localstore.New(localstore.WithClock(func() time.Time { return fixed }))
|
|
template := templates.Template{ID: "tpl", OwnerID: "owner", Name: "T", Prompt: "P", CreatedAt: fixed, UpdatedAt: fixed}
|
|
if _, err := store.CreateTemplate(ctx, template); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if items, _ := store.ListTemplates(ctx, "other"); len(items) != 0 {
|
|
t.Fatal("template leaked")
|
|
}
|
|
event := usage.Event{ID: "usage-1", JobID: "job-1", OwnerID: "owner", Capability: "image.generate", CreatedAt: fixed.Format(time.RFC3339)}
|
|
if _, inserted, _ := store.Insert(event); !inserted {
|
|
t.Fatal("first usage not inserted")
|
|
}
|
|
if _, inserted, _ := store.Insert(event); inserted {
|
|
t.Fatal("duplicate job usage inserted")
|
|
}
|
|
rules := billing.DefaultBillingPriceRules()[:1]
|
|
if err := store.SeedBillingPriceRules(ctx, rules); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.SeedBillingPriceRules(ctx, rules); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, _ := store.ListBillingPriceRules(ctx, true)
|
|
if len(got) != 1 {
|
|
t.Fatalf("rules = %d", len(got))
|
|
}
|
|
}
|