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