geo/server/internal/tenant/app/ai_point_usage_integration_test.go

//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)
	}
}