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

package app

import (
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
)

func TestCalculateAIPointCost(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		chars     int
		baseChars int
		want      int
	}{
		{name: "empty still charges base usage", chars: 0, baseChars: 1000, want: 1},
		{name: "single char", chars: 1, baseChars: 1000, want: 1},
		{name: "first base boundary", chars: 1000, baseChars: 1000, want: 1},
		{name: "over first base", chars: 1001, baseChars: 1000, want: 2},
		{name: "second base boundary", chars: 2000, baseChars: 1000, want: 2},
		{name: "over second base", chars: 2001, baseChars: 1000, want: 3},
		{name: "default base chars", chars: 1000, baseChars: 0, want: 1},
		{name: "custom base chars", chars: 1500, baseChars: 500, want: 3},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			if got := calculateAIPointCost(tt.chars, tt.baseChars); got != tt.want {
				t.Fatalf("calculateAIPointCost(%d, %d) = %d, want %d", tt.chars, tt.baseChars, got, tt.want)
			}
		})
	}
}

func TestCalculateAIPointOutputSettlementUsesOutputText(t *testing.T) {
	t.Parallel()

	reservation := AIPointReservation{
		Points:       1,
		RequestChars: 8,
		BaseChars:    1000,
	}

	got := calculateAIPointOutputSettlement(reservation, strings.Repeat("贵", 2500))

	if got.OutputChars != 2500 {
		t.Fatalf("OutputChars = %d, want 2500", got.OutputChars)
	}
	if got.FinalPoints != 3 {
		t.Fatalf("FinalPoints = %d, want 3", got.FinalPoints)
	}
	if got.LedgerDelta != -2 {
		t.Fatalf("LedgerDelta = %d, want -2 for output-side top-up", got.LedgerDelta)
	}
}

func TestCalculateAIPointOutputSettlementBoundaries(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name        string
		outputChars int
		wantPoints  int
	}{
		{name: "minimum empty output still one point", outputChars: 0, wantPoints: 1},
		{name: "below base", outputChars: 999, wantPoints: 1},
		{name: "exactly base", outputChars: 1000, wantPoints: 1},
		{name: "one over base", outputChars: 1001, wantPoints: 2},
		{name: "exactly second base", outputChars: 2000, wantPoints: 2},
		{name: "one over second base", outputChars: 2001, wantPoints: 3},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			got := calculateAIPointOutputSettlement(
				AIPointReservation{Points: 1, RequestChars: 50000, BaseChars: 1000},
				strings.Repeat("字", tt.outputChars),
			)
			if got.FinalPoints != tt.wantPoints {
				t.Fatalf("FinalPoints = %d, want %d", got.FinalPoints, tt.wantPoints)
			}
		})
	}
}

func TestCalculateAIPointOutputSettlementRefundsOverReservation(t *testing.T) {
	t.Parallel()

	got := calculateAIPointOutputSettlement(
		AIPointReservation{Points: 5, RequestChars: 5000, BaseChars: 1000},
		strings.Repeat("字", 700),
	)

	if got.FinalPoints != 1 {
		t.Fatalf("FinalPoints = %d, want 1", got.FinalPoints)
	}
	if got.LedgerDelta != 4 {
		t.Fatalf("LedgerDelta = %d, want +4 for output-side refund", got.LedgerDelta)
	}
}

func TestCalculateAIPointOutputSettlementUsesDefaultBaseChars(t *testing.T) {
	t.Parallel()

	got := calculateAIPointOutputSettlement(
		AIPointReservation{Points: 1, BaseChars: 0},
		strings.Repeat("字", 1001),
	)

	if got.BaseChars != 1000 {
		t.Fatalf("BaseChars = %d, want default 1000", got.BaseChars)
	}
	if got.FinalPoints != 2 {
		t.Fatalf("FinalPoints = %d, want 2", got.FinalPoints)
	}
	if got.LedgerDelta != -1 {
		t.Fatalf("LedgerDelta = %d, want -1", got.LedgerDelta)
	}
}

func TestCompleteAIPointReservationWithFinalAmountUpdatesReservationAmounts(t *testing.T) {
	tx := &captureAIPointTx{}

	adjusted, err := completeAIPointReservationWithFinalAmountIfPending(context.Background(), tx, 10, 20, 3)
	if err != nil {
		t.Fatalf("complete reservation: %v", err)
	}
	if !adjusted {
		t.Fatalf("adjusted = false, want true")
	}
	if !strings.Contains(tx.sql, "reserved_amount = $4") {
		t.Fatalf("reservation SQL must set final reserved_amount; query:\n%s", tx.sql)
	}
	if !strings.Contains(tx.sql, "consumed_amount = $4") {
		t.Fatalf("reservation SQL must set final consumed_amount; query:\n%s", tx.sql)
	}
	if len(tx.args) != 4 {
		t.Fatalf("args length = %d, want 4", len(tx.args))
	}
	if tx.args[0] != int64(20) || tx.args[1] != int64(10) || tx.args[2] != aiPointsQuotaType || tx.args[3] != 3 {
		t.Fatalf("args = %#v, want reservation, tenant, quota type, final points", tx.args)
	}
}

type captureAIPointTx struct {
	sql  string
	args []any
}

func (tx *captureAIPointTx) Begin(context.Context) (pgx.Tx, error) { return nil, errors.New("unused") }
func (tx *captureAIPointTx) Commit(context.Context) error          { return nil }
func (tx *captureAIPointTx) Rollback(context.Context) error        { return nil }
func (tx *captureAIPointTx) CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) {
	return 0, errors.New("unused")
}
func (tx *captureAIPointTx) SendBatch(context.Context, *pgx.Batch) pgx.BatchResults {
	return nil
}
func (tx *captureAIPointTx) LargeObjects() pgx.LargeObjects { return pgx.LargeObjects{} }
func (tx *captureAIPointTx) Prepare(context.Context, string, string) (*pgconn.StatementDescription, error) {
	return nil, errors.New("unused")
}
func (tx *captureAIPointTx) Exec(_ context.Context, sql string, args ...any) (pgconn.CommandTag, error) {
	tx.sql = sql
	tx.args = args
	return pgconn.NewCommandTag("UPDATE 1"), nil
}
func (tx *captureAIPointTx) Query(context.Context, string, ...any) (pgx.Rows, error) {
	return nil, errors.New("unused")
}
func (tx *captureAIPointTx) QueryRow(context.Context, string, ...any) pgx.Row {
	return failingAIPointRow{}
}
func (tx *captureAIPointTx) Conn() *pgx.Conn { return nil }

type failingAIPointRow struct{}

func (failingAIPointRow) Scan(...any) error { return errors.New("unused") }