descrybe/apps/api/internal/billing/consume_credits_integration_test.go

package billing

import (
	"context"
	"errors"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Concurrent ConsumeCredits on one company must serialize on credit_balances and
// never overspend the wallet.
func TestConsumeCreditsConcurrentNoOverspend(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, "consume-credits-contention")
	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)
	})

	// Paid plan keeps CanUseAI true at empty wallet so flat (0-token) debits still
	// hit the atomic UPDATE and return ErrInsufficientCredits (not a Free no-op).
	var planID int64
	err = pg.QueryRow(ctx, `
		INSERT INTO plans (name, description, monthly_credits, term)
		VALUES ($1, $2, $3, 'monthly')
		RETURNING id`, "consume-contention-"+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)
	})
	_, err = pg.Exec(ctx, `
		INSERT INTO company_plans (company_id, plan_id, is_active, billing_cycle_start, next_billing_date)
		VALUES ($1, $2, true, now() - interval '1 day', now() + interval '30 days')`, companyID, planID)
	if err != nil {
		t.Fatal(err)
	}

	const wallet = 20
	_, err = pg.Exec(ctx, `
		INSERT INTO credit_balances (company_id, total_credits, used_credits, updated_at)
		VALUES ($1, $2, 0, now())`, companyID, wallet)
	if err != nil {
		t.Fatal(err)
	}
	_, err = pg.Exec(ctx, `
		INSERT INTO billing_cycles (company_id, start_date, end_date, credits_used, products_processed)
		VALUES ($1, now() - interval '1 day', now() + interval '30 days', 0, 0)`, companyID)
	if err != nil {
		t.Fatal(err)
	}

	svc := &Service{Pool: pg}
	_ = svc.EnsureDefaultCosts(ctx)

	const workers = 40
	var wg sync.WaitGroup
	var okCount atomic.Int64
	var insuff atomic.Int64
	startGate := make(chan struct{})

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-startGate
			err := svc.ConsumeCredits(ctx, companyID, 0, "product_processing")
			if err == nil {
				okCount.Add(1)
				return
			}
			if errors.Is(err, ErrInsufficientCredits) {
				insuff.Add(1)
				return
			}
			t.Errorf("unexpected: %v", err)
		}()
	}
	close(startGate)
	wg.Wait()

	if okCount.Load() != wallet {
		t.Fatalf("ok=%d want %d (insuff=%d)", okCount.Load(), wallet, insuff.Load())
	}
	if okCount.Load()+insuff.Load() != workers {
		t.Fatalf("ok+insuff=%d want %d", okCount.Load()+insuff.Load(), workers)
	}

	var used, total 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 used != wallet || total != wallet {
		t.Fatalf("wallet total=%d used=%d want total=%d used=%d", total, used, wallet, wallet)
	}

	var cycleUsed, products int
	err = pg.QueryRow(ctx, `
		SELECT credits_used, products_processed FROM billing_cycles
		WHERE company_id = $1 AND end_date > now()
		ORDER BY start_date DESC LIMIT 1`, companyID).Scan(&cycleUsed, &products)
	if err != nil {
		t.Fatal(err)
	}
	if cycleUsed != wallet || products != wallet {
		t.Fatalf("cycle used=%d products=%d want %d", cycleUsed, products, wallet)
	}
}

func TestConsumeCreditsBatchMatchesSummedBase(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, "consume-credits-batch")
	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)
	})

	_, err = pg.Exec(ctx, `
		INSERT INTO credit_balances (company_id, total_credits, used_credits, updated_at)
		VALUES ($1, 100, 0, now())`, companyID)
	if err != nil {
		t.Fatal(err)
	}
	_, err = pg.Exec(ctx, `
		INSERT INTO billing_cycles (company_id, start_date, end_date, credits_used, products_processed)
		VALUES ($1, now() - interval '1 day', now() + interval '30 days', 0, 0)`, companyID)
	if err != nil {
		t.Fatal(err)
	}

	svc := &Service{Pool: pg}
	_ = svc.EnsureDefaultCosts(ctx)

	// 5 products, 0 tokens → DebitAmountN = 5 base credits, products_processed += 5.
	if err := svc.ConsumeCreditsBatch(ctx, companyID, 0, 5, "product_processing"); err != nil {
		t.Fatal(err)
	}

	var used, products int
	err = pg.QueryRow(ctx, `
		SELECT cb.used_credits, bc.products_processed
		FROM credit_balances cb
		JOIN billing_cycles bc ON bc.company_id = cb.company_id AND bc.end_date > now()
		WHERE cb.company_id = $1
		ORDER BY bc.start_date DESC LIMIT 1`, companyID).Scan(&used, &products)
	if err != nil {
		t.Fatal(err)
	}
	if used != 5 || products != 5 {
		t.Fatalf("used=%d products=%d want 5/5", used, products)
	}
}