package billing import ( "context" "os" "sync" "testing" "time" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" ) // Concurrent claimAndRollDueCompanyPlan must roll a due company_plan exactly once. func TestClaimAndRollDueCompanyPlanConcurrent(t *testing.T) { dsn := os.Getenv("DATABASE_URL") if dsn == "" { t.Skip("DATABASE_URL not set") } ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() pg, err := pgxpool.New(ctx, dsn) if err != nil { t.Fatal(err) } defer pg.Close() companyID := uuid.New() _, err = pg.Exec(ctx, `INSERT INTO companies (id, name) VALUES ($1, $2)`, companyID, "billing-cycle-claim-test") if err != nil { t.Fatal(err) } t.Cleanup(func() { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) defer cleanupCancel() _, _ = pg.Exec(cleanupCtx, `DELETE FROM companies WHERE id = $1`, companyID) }) var planID int64 err = pg.QueryRow(ctx, ` INSERT INTO plans (name, description, monthly_credits, term) VALUES ($1, $2, $3, 'monthly') RETURNING id`, "claim-test-plan-"+companyID.String()[:8], "integration", 100).Scan(&planID) if err != nil { t.Fatal(err) } t.Cleanup(func() { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) defer cleanupCancel() _, _ = pg.Exec(cleanupCtx, `DELETE FROM plans WHERE id = $1`, planID) }) cycleStart := time.Now().UTC().AddDate(0, -1, 0) nextBill := time.Now().UTC().Add(-time.Hour) var rowID int64 err = pg.QueryRow(ctx, ` INSERT INTO company_plans (company_id, plan_id, is_active, billing_cycle_start, next_billing_date) VALUES ($1, $2, true, $3, $4) RETURNING id`, companyID, planID, cycleStart, nextBill).Scan(&rowID) if err != nil { t.Fatal(err) } _, err = pg.Exec(ctx, ` INSERT INTO credit_balances (company_id, total_credits, used_credits, updated_at) VALUES ($1, 100, 17, now())`, companyID) if err != nil { t.Fatal(err) } svc := &Service{Pool: pg} const workers = 8 var wg sync.WaitGroup errs := make(chan error, workers) oks := make(chan bool, workers) startGate := make(chan struct{}) for i := 0; i < workers; i++ { wg.Add(1) go func() { defer wg.Done() <-startGate ok, runErr := svc.claimAndRollDueCompanyPlan(ctx, rowID) if runErr != nil { errs <- runErr return } oks <- ok }() } close(startGate) wg.Wait() close(errs) close(oks) for err := range errs { t.Fatalf("claimAndRollDueCompanyPlan: %v", err) } successes := 0 for ok := range oks { if ok { successes++ } } if successes != 1 { t.Fatalf("expected exactly 1 successful claim, got %d", successes) } var cycleCount int err = pg.QueryRow(ctx, `SELECT COUNT(*) FROM billing_cycles WHERE company_id = $1`, companyID).Scan(&cycleCount) if err != nil { t.Fatal(err) } if cycleCount != 1 { t.Fatalf("billing_cycles rows=%d want 1", cycleCount) } var creditsUsed int err = pg.QueryRow(ctx, ` SELECT credits_used FROM billing_cycles WHERE company_id = $1`, companyID).Scan(&creditsUsed) if err != nil { t.Fatal(err) } if creditsUsed != 17 { t.Fatalf("credits_used=%d want 17", creditsUsed) } var stillDue bool err = pg.QueryRow(ctx, ` SELECT next_billing_date <= now() FROM company_plans WHERE id = $1`, rowID).Scan(&stillDue) if err != nil { t.Fatal(err) } if stillDue { t.Fatal("company_plans still due after roll") } var total, used int err = pg.QueryRow(ctx, ` SELECT total_credits, used_credits FROM credit_balances WHERE company_id = $1`, companyID).Scan(&total, &used) if err != nil { t.Fatal(err) } if total != 100 || used != 0 { t.Fatalf("credit_balances total=%d used=%d want 100/0", total, used) } } func TestRunDueBillingCyclesBestEffortMultiCompany(t *testing.T) { dsn := os.Getenv("DATABASE_URL") if dsn == "" { t.Skip("DATABASE_URL not set") } ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() pg, err := pgxpool.New(ctx, dsn) if err != nil { t.Fatal(err) } defer pg.Close() type fixture struct { companyID uuid.UUID planID int64 rowID int64 } var fixtures []fixture for i := 0; i < 2; i++ { companyID := uuid.New() _, err = pg.Exec(ctx, `INSERT INTO companies (id, name) VALUES ($1, $2)`, companyID, "billing-cycle-multi-"+companyID.String()[:8]) if err != nil { t.Fatal(err) } cid := companyID t.Cleanup(func() { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) defer cleanupCancel() _, _ = pg.Exec(cleanupCtx, `DELETE FROM companies WHERE id = $1`, cid) }) var planID int64 err = pg.QueryRow(ctx, ` INSERT INTO plans (name, description, monthly_credits, term) VALUES ($1, $2, $3, 'monthly') RETURNING id`, "multi-plan-"+companyID.String()[:8], "integration", 80).Scan(&planID) if err != nil { t.Fatal(err) } pid := planID t.Cleanup(func() { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) defer cleanupCancel() _, _ = pg.Exec(cleanupCtx, `DELETE FROM plans WHERE id = $1`, pid) }) cycleStart := time.Now().UTC().AddDate(0, -1, 0) nextBill := time.Now().UTC().Add(-time.Hour) var rowID int64 err = pg.QueryRow(ctx, ` INSERT INTO company_plans (company_id, plan_id, is_active, billing_cycle_start, next_billing_date) VALUES ($1, $2, true, $3, $4) RETURNING id`, companyID, planID, cycleStart, nextBill).Scan(&rowID) if err != nil { t.Fatal(err) } _, err = pg.Exec(ctx, ` INSERT INTO credit_balances (company_id, total_credits, used_credits, updated_at) VALUES ($1, 80, 3, now())`, companyID) if err != nil { t.Fatal(err) } fixtures = append(fixtures, fixture{companyID: companyID, planID: planID, rowID: rowID}) } svc := &Service{Pool: pg} res, runErr := svc.RunDueBillingCycles(ctx) if runErr != nil { t.Fatalf("RunDueBillingCycles: %v", runErr) } if res.Processed < 2 { t.Fatalf("processed=%d want >=2 (got failed=%d)", res.Processed, res.Failed) } if res.Failed != 0 { t.Fatalf("failed=%d want 0", res.Failed) } for _, f := range fixtures { var stillDue bool err = pg.QueryRow(ctx, ` SELECT next_billing_date <= now() FROM company_plans WHERE id = $1`, f.rowID).Scan(&stillDue) if err != nil { t.Fatal(err) } if stillDue { t.Fatalf("company_plan %d still due after multi-company run", f.rowID) } } }