fix(growth): harden AI jobs and cost controls

internal/growthreport/repository_integration_test.go

@@ -0,0 +1,81 @@
+package growthreport
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func TestRepositoryRecoversExpiredActiveLeases(t *testing.T) {
+	pool, merchantID := growthRepositoryTestPool(t)
+	ctx := context.Background()
+	now := time.Now().UTC()
+	for index, status := range []Status{StatusAnalyzing, StatusGenerating} {
+		id := fmt.Sprintf("growth_%024x", index+1)
+		_, err := pool.Exec(ctx, `INSERT INTO growth_reports(id,merchant_id,fingerprint,status,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,lease_token,lease_expires_at,created_at,updated_at) VALUES($1,$2,$3,$4,current_date,'test','test','test','test','test','test',$5,$6,$7,$7)`, id, merchantID, id, status, "expired-token", now.Add(-time.Minute), now.Add(time.Duration(index)*time.Second))
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	repository := NewRepository(pool)
+	first, err := repository.Claim(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if first.Status != StatusAnalyzing || first.ClaimToken == "" || first.ClaimToken == "expired-token" {
+		t.Fatalf("recovered claim=%+v", first)
+	}
+	if err := repository.CompleteClaimed(ctx, first.ReportID, "expired-token", StatusPartial, json.RawMessage(`{}`), nil, Usage{}, "stale"); !errors.Is(err, ErrNotFound) {
+		t.Fatalf("stale completion error=%v", err)
+	}
+	second, err := repository.Claim(ctx)
+	if err != nil || second.ReportID == first.ReportID || second.ClaimToken == "" {
+		t.Fatalf("second claim=%+v error=%v", second, err)
+	}
+}
+
+func TestRepositoryChecksUncachedAllowanceBeforeInsert(t *testing.T) {
+	pool, merchantID := growthRepositoryTestPool(t)
+	repository := NewRepository(pool)
+	start := time.Now().UTC()
+	input := CreateRecord{ID: "growth_ffffffffffffffffffffffff", MerchantID: merchantID, Fingerprint: "blocked", DatasetWatermark: "test", ForecastStartDate: &start, CatalogVersion: "test", AnalyticsVersion: "test", PromptVersion: "test", SchemaVersion: "test", Model: "test"}
+	_, _, err := repository.CreateOrReuse(context.Background(), input, func() error { return ErrRateLimited })
+	if !errors.Is(err, ErrRateLimited) {
+		t.Fatalf("CreateOrReuse error=%v", err)
+	}
+	var count int
+	if err := pool.QueryRow(context.Background(), `SELECT count(*) FROM growth_reports WHERE id=$1`, input.ID).Scan(&count); err != nil || count != 0 {
+		t.Fatalf("queued reports=%d error=%v", count, err)
+	}
+}
+
+func growthRepositoryTestPool(t *testing.T) (*pgxpool.Pool, int64) {
+	t.Helper()
+	databaseURL := os.Getenv("TEST_DATABASE_URL")
+	if databaseURL == "" {
+		t.Skip("TEST_DATABASE_URL is not configured")
+	}
+	pool, err := pgxpool.New(context.Background(), databaseURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sourceKey := fmt.Sprintf("growth-test-%d", time.Now().UnixNano())
+	var merchantID int64
+	if err := pool.QueryRow(context.Background(), `INSERT INTO merchants(source_key,category_id,category_title,display_name) VALUES($1,1,'test','test') RETURNING id`, sourceKey).Scan(&merchantID); err != nil {
+		pool.Close()
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		_, _ = pool.Exec(context.Background(), `DELETE FROM growth_reports WHERE merchant_id=$1`, merchantID)
+		_, _ = pool.Exec(context.Background(), `DELETE FROM merchants WHERE id=$1`, merchantID)
+		pool.Close()
+	})
+	return pool, merchantID
+}