package main import ( "os" "path/filepath" "strconv" "sync" "testing" "github.com/glebarez/sqlite" "gorm.io/gorm" "time" ) func TestClaimTaskIsAtomic(t *testing.T) { a := newCreditTestApp(t) task := Task{UserID: 1, Source: "kw", RID: "atomic", Quality: "lossless", Status: "pending"} if err := a.db.Create(&task).Error; err != nil { t.Fatal(err) } if _, err := a.claimTask(task.ID); err != nil { t.Fatal(err) } if _, err := a.claimTask(task.ID); err != errTaskRunning { t.Fatalf("second claim error = %v, want task running", err) } } func TestRetryCommittedReservationDoesNotChargeAgain(t *testing.T) { a := newCreditTestApp(t) user := User{Username: "retry-committed", Role: RoleUser, Status: UserActive} if err := a.db.Create(&user).Error; err != nil { t.Fatal(err) } if err := a.db.Create(&CreditAccount{UserID: user.ID, Balance: 1}).Error; err != nil { t.Fatal(err) } task := Task{UserID: user.ID, Source: "kw", RID: "committed", Quality: "lossless", Status: "pending"} if err := a.db.Create(&task).Error; err != nil { t.Fatal(err) } r, err := a.reserveCredit(user.ID, task.ID, 1) if err != nil { t.Fatal(err) } task.CreditReservationID = r.ID if err := a.db.Model(&task).Update("credit_reservation_id", r.ID).Error; err != nil { t.Fatal(err) } if err := a.commitCredit(&task); err != nil { t.Fatal(err) } if err := a.db.Model(&task).Updates(map[string]any{"status": "error", "error": "legacy"}).Error; err != nil { t.Fatal(err) } var current Task if err := a.db.First(¤t, task.ID).Error; err != nil { t.Fatal(err) } if err := a.prepareRetry(¤t); err != nil { t.Fatal(err) } var account CreditAccount if err := a.db.Where("user_id = ?", user.ID).First(&account).Error; err != nil { t.Fatal(err) } if account.Balance != 0 || account.Reserved != 0 { t.Fatalf("credit changed on committed retry: %+v", account) } if current.Status != "upload_pending" { t.Fatalf("status = %q, want upload_pending", current.Status) } } func TestRetryRefundedReservationCreatesNewReservation(t *testing.T) { a := newCreditTestApp(t) user := User{Username: "retry-refunded", Role: RoleUser, Status: UserActive} if err := a.db.Create(&user).Error; err != nil { t.Fatal(err) } if err := a.db.Create(&CreditAccount{UserID: user.ID, Balance: 1}).Error; err != nil { t.Fatal(err) } task := Task{UserID: user.ID, Source: "kw", RID: "refunded", Quality: "lossless", Status: "error"} if err := a.db.Create(&task).Error; err != nil { t.Fatal(err) } r1, err := a.reserveCredit(user.ID, task.ID, 1) if err != nil { t.Fatal(err) } task.CreditReservationID = r1.ID if err := a.refundCredit(&task); err != nil { t.Fatal(err) } if err := a.prepareRetry(&task); err != nil { t.Fatal(err) } if task.CreditReservationID == r1.ID { t.Fatal("retry reused refunded reservation") } var account CreditAccount if err := a.db.Where("user_id = ?", user.ID).First(&account).Error; err != nil { t.Fatal(err) } if account.Balance != 1 || account.Reserved != 1 { t.Fatalf("unexpected account: %+v", account) } } func TestStaleAssetLeaseStealsSafely(t *testing.T) { a := newCreditTestApp(t) root := t.TempDir() a.cfg.DownloadDir = root first := Task{UserID: 1, Source: "kw", RID: "lease", Quality: "lossless", Format: "flac", Status: "downloading"} second := Task{UserID: 2, Source: "kw", RID: "lease", Quality: "lossless", Format: "flac", Status: "pending"} if err := a.db.Create(&first).Error; err != nil { t.Fatal(err) } if err := a.db.Create(&second).Error; err != nil { t.Fatal(err) } asset := MediaAsset{ AssetKey: assetKey("kw", "lease", "lossless"), Source: "kw", RID: "lease", Quality: "lossless", LocalPath: filepath.Join(root, "old.flac"), Format: "flac", Status: "downloading", OwnerTaskID: first.ID, LeaseUntil: time.Now().Add(-time.Minute), Generation: 1, } if err := a.db.Create(&asset).Error; err != nil { t.Fatal(err) } first.AssetID = asset.ID if err := a.db.Model(&first).Update("asset_id", asset.ID).Error; err != nil { t.Fatal(err) } got, reused, err := a.acquireAsset(&second) if err != nil || reused { t.Fatalf("acquire = asset=%+v reused=%v err=%v", got, reused, err) } if got.Generation != 2 { t.Fatalf("generation = %d, want 2", got.Generation) } if got.LocalPath == asset.LocalPath { t.Fatal("stolen generation reused old local path") } var old Task if err := a.db.First(&old, first.ID).Error; err != nil { t.Fatal(err) } if old.Status != "pending" { t.Fatalf("stale owner status = %q, want pending", old.Status) } } func TestAssetCachePathIsStablePerGeneration(t *testing.T) { root := t.TempDir() p1 := assetCachePath(root, "kw-rid-lossless", "flac", 1) p2 := assetCachePath(root, "kw-rid-lossless", "flac", 2) if p1 == p2 || filepath.Dir(p1) != filepath.Dir(p2) { t.Fatalf("unexpected cache paths: %q %q", p1, p2) } if err := os.MkdirAll(filepath.Dir(p1), 0755); err != nil { t.Fatal(err) } } func TestConcurrentCreditReservationsAreAtomic(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "credits.db") db, err := gorm.Open(sqlite.Open(dbPath+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(10000)"), &gorm.Config{}) if err != nil { t.Fatal(err) } if err := db.AutoMigrate(&Task{}, &MediaAsset{}, &User{}, &CreditAccount{}, &CreditTransaction{}, &CreditReservation{}); err != nil { t.Fatal(err) } app := &App{db: db} user := User{Username: "concurrent-credit", Role: RoleUser, Status: UserActive} if err := db.Create(&user).Error; err != nil { t.Fatal(err) } if err := db.Create(&CreditAccount{UserID: user.ID, Balance: 5}).Error; err != nil { t.Fatal(err) } for i := 0; i < 20; i++ { if err := db.Create(&Task{UserID: user.ID, Source: "kw", RID: "c-" + strconv.Itoa(i), Quality: "lossless", Status: "pending"}).Error; err != nil { t.Fatal(err) } } var wg sync.WaitGroup var mu sync.Mutex success, unexpected := 0, 0 for i := 0; i < 20; i++ { var task Task if err := db.Where("user_id = ? AND r_id = ?", user.ID, "c-"+strconv.Itoa(i)).First(&task).Error; err != nil { t.Fatal(err) } wg.Add(1) go func(taskID uint) { defer wg.Done() if _, err := app.reserveCredit(user.ID, taskID, 1); err == nil { mu.Lock() success++ mu.Unlock() } else if err != errInsufficientCredits { mu.Lock() unexpected++ mu.Unlock() } }(task.ID) } wg.Wait() if success != 5 || unexpected != 0 { t.Fatalf("reservations: success=%d unexpected=%d", success, unexpected) } var account CreditAccount if err := db.Where("user_id = ?", user.ID).First(&account).Error; err != nil { t.Fatal(err) } if account.Balance != 5 || account.Reserved != 5 { t.Fatalf("account after concurrent reservations: %+v", account) } }