//go:build integration package app import ( "context" "fmt" "os" "strings" "testing" "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/geo-platform/tenant-api/internal/tenant/repository" ) func TestAIPointsOutputSettlementRealDatabase(t *testing.T) { dsn := os.Getenv("TEST_DATABASE_URL") if dsn == "" { t.Skip("set TEST_DATABASE_URL to run the real database AI point settlement test") } ctx := context.Background() pool, err := pgxpool.New(ctx, dsn) if err != nil { t.Fatalf("connect database: %v", err) } t.Cleanup(pool.Close) tenantID, userID := seedAIPointSettlementTenant(t, ctx, pool, 10, 1000) reservation, err := ReserveAIPoints(ctx, pool, nil, AIPointReserveInput{ TenantID: tenantID, OperatorID: userID, UsageType: "test_output_billing", MeteredText: "cheap input", }) if err != nil { t.Fatalf("reserve ai points: %v", err) } if reservation.Points != 1 { t.Fatalf("reserved points = %d, want minimum 1", reservation.Points) } output := strings.Repeat("่ดต", 2500) completed, err := CompleteAIPoints(ctx, pool, nil, tenantID, *reservation, "test-model", output) if err != nil { t.Fatalf("complete ai points: %v", err) } if completed.Points != 3 { t.Fatalf("completed points = %d, want 3", completed.Points) } if completed.RequestChars != 2500 { t.Fatalf("completed request/output chars = %d, want 2500", completed.RequestChars) } var usageStatus, model string var usageChars, usagePoints int if err := pool.QueryRow(ctx, ` SELECT status, request_chars, points, model FROM ai_point_usage_logs WHERE id = $1 AND tenant_id = $2 `, reservation.UsageLogID, tenantID).Scan(&usageStatus, &usageChars, &usagePoints, &model); err != nil { t.Fatalf("query usage log: %v", err) } if usageStatus != "completed" || usageChars != 2500 || usagePoints != 3 || model != "test-model" { t.Fatalf("usage log = status %q chars %d points %d model %q; want completed 2500 3 test-model", usageStatus, usageChars, usagePoints, model) } var reservationStatus string var reservedAmount, consumedAmount int if err := pool.QueryRow(ctx, ` SELECT status, reserved_amount, consumed_amount FROM quota_reservations WHERE id = $1 AND tenant_id = $2 `, reservation.ReservationID, tenantID).Scan(&reservationStatus, &reservedAmount, &consumedAmount); err != nil { t.Fatalf("query quota reservation: %v", err) } if reservationStatus != "confirmed" || reservedAmount != 3 || consumedAmount != 3 { t.Fatalf("reservation = status %q reserved %d consumed %d; want confirmed 3 3", reservationStatus, reservedAmount, consumedAmount) } rows, err := pool.Query(ctx, ` SELECT delta, balance_after, reason FROM tenant_quota_ledgers WHERE tenant_id = $1 AND quota_type = 'ai_points' AND reference_id = $2 ORDER BY id `, tenantID, reservation.UsageLogID) if err != nil { t.Fatalf("query quota ledger: %v", err) } defer rows.Close() var deltas []int var balances []int var reasons []string for rows.Next() { var delta, balance int var reason string if err := rows.Scan(&delta, &balance, &reason); err != nil { t.Fatalf("scan quota ledger: %v", err) } deltas = append(deltas, delta) balances = append(balances, balance) reasons = append(reasons, reason) } if err := rows.Err(); err != nil { t.Fatalf("iterate quota ledger: %v", err) } if fmt.Sprint(deltas) != "[-1 -2]" || fmt.Sprint(balances) != "[9 7]" || fmt.Sprint(reasons) != "[ai_points_reserve ai_points_output_settle]" { t.Fatalf("ledger deltas=%v balances=%v reasons=%v; want [-1 -2], [9 7], reserve+settle", deltas, balances, reasons) } status, err := repository.NewQuotaRepository(pool).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) if err != nil { t.Fatalf("query quota status: %v", err) } if status.Used != 3 || status.Balance != 7 { t.Fatalf("quota status used=%d balance=%d; want 3/7", status.Used, status.Balance) } } func seedAIPointSettlementTenant(t *testing.T, ctx context.Context, pool *pgxpool.Pool, aiPointsMonthly int, baseChars int) (int64, int64) { t.Helper() suffix := fmt.Sprintf("%d", time.Now().UnixNano()) var tenantID int64 if err := pool.QueryRow(ctx, ` INSERT INTO tenants (name, status) VALUES ($1, 'active') RETURNING id `, "ai-point-output-test-"+suffix).Scan(&tenantID); err != nil { t.Fatalf("insert tenant: %v", err) } var userID int64 if err := pool.QueryRow(ctx, ` INSERT INTO users (phone, password_hash, name, status) VALUES ($1, 'test', 'AI Point Output Test', 'active') RETURNING id `, "199"+suffix[len(suffix)-8:]).Scan(&userID); err != nil { t.Fatalf("insert user: %v", err) } if _, err := pool.Exec(ctx, ` INSERT INTO tenant_memberships (tenant_id, user_id, tenant_role) VALUES ($1, $2, 'owner') `, tenantID, userID); err != nil { t.Fatalf("insert tenant membership: %v", err) } var planID int64 if err := pool.QueryRow(ctx, ` INSERT INTO plans (plan_code, name, quota_policy_json, status) VALUES ($1, 'AI Point Output Plan', jsonb_build_object( 'article_generation', 0, 'article_quota_cycle', 'monthly', 'ai_points_monthly', $2::int, 'ai_point_base_chars', $3::int, 'image_storage_bytes', 0, 'brand_limit', 1 ), 'active') RETURNING id `, "ai-point-output-"+suffix, aiPointsMonthly, baseChars).Scan(&planID); err != nil { t.Fatalf("insert plan: %v", err) } if _, err := pool.Exec(ctx, ` INSERT INTO tenant_plan_subscriptions (tenant_id, plan_id, start_at, end_at, status) VALUES ($1, $2, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '24 hours', 'active') `, tenantID, planID); err != nil { t.Fatalf("insert tenant plan subscription: %v", err) } t.Cleanup(func() { cleanupCtx := context.Background() mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM ai_point_usage_logs WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_quota_ledgers WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM quota_reservations WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_plan_subscriptions WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspace_memberships WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM workspaces WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenant_memberships WHERE tenant_id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM users WHERE id = $1`, userID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM tenants WHERE id = $1`, tenantID) mustCleanupAIPointSettlementTestData(t, cleanupCtx, pool, `DELETE FROM plans WHERE id = $1`, planID) }) return tenantID, userID } func mustCleanupAIPointSettlementTestData(t *testing.T, ctx context.Context, pool *pgxpool.Pool, sql string, args ...any) { t.Helper() if _, err := pool.Exec(ctx, sql, args...); err != nil { t.Errorf("cleanup AI point settlement test data: %v; sql=%s", err, sql) } }