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

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