186 lines
8.7 KiB
Go
186 lines
8.7 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
)
|
|
|
|
func TestChargeAndActivateCreationCommitsWalletAndDispatchGateTogether(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
tx := &generationStateTransaction{results: []*jobRows{
|
|
{rows: [][]any{{"ledger-1", int64(965), int64(965), int64(1000), int64(35), now, now, int64(-35)}}},
|
|
{rows: [][]any{{"job-1"}}},
|
|
}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, &generationStatePool{tx: tx})
|
|
got, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", AccountID: "account", JobID: "job-1", AmountFen: 35, Description: "image"}, json.RawMessage(`{"status":"pending","amountFen":35}`))
|
|
if err != nil || tx.commits != 1 || tx.rollbacks != 0 || len(tx.queries) != 2 {
|
|
t.Fatalf("got=%s err=%v commits=%d rollbacks=%d queries=%d", got, err, tx.commits, tx.rollbacks, len(tx.queries))
|
|
}
|
|
if tx.queries[0].sql != PostWalletEntrySQL || tx.queries[1].sql != activateChargedCreationSQL {
|
|
t.Fatalf("queries=%#v", tx.queries)
|
|
}
|
|
if tx.queries[0].args[7] != "job-charge:job-1" || tx.queries[1].args[0] != "job-1" {
|
|
t.Fatalf("args=%#v", tx.queries)
|
|
}
|
|
var state map[string]any
|
|
_ = json.Unmarshal(got, &state)
|
|
if state["status"] != "charged" || state["ledgerEntryId"] != "ledger-1" {
|
|
t.Fatalf("billing=%#v", state)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationRollsBackWhenActivationFails(t *testing.T) {
|
|
now := time.Now()
|
|
tx := &generationStateTransaction{results: []*jobRows{
|
|
{rows: [][]any{{"ledger-1", int64(1), int64(1), int64(1), int64(1), now, now, int64(-1)}}},
|
|
{},
|
|
}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, &generationStatePool{tx: tx})
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job", AmountFen: 1}, json.RawMessage(`{"status":"pending"}`))
|
|
if err == nil || tx.commits != 0 || tx.rollbacks != 1 {
|
|
t.Fatalf("err=%v commits=%d rollbacks=%d", err, tx.commits, tx.rollbacks)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationPreservesDeferredWalletError(t *testing.T) {
|
|
deferred := &pgconn.PgError{Code: "P0001", Message: "BILLING_INSUFFICIENT_BALANCE"}
|
|
tx := &generationStateTransaction{results: []*jobRows{{err: deferred}}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, &generationStatePool{tx: tx})
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job", AmountFen: 39}, json.RawMessage(`{"status":"pending"}`))
|
|
var postgresErr *pgconn.PgError
|
|
if !errors.As(err, &postgresErr) || postgresErr.Code != "P0001" || tx.rollbacks != 1 {
|
|
t.Fatalf("err=%v postgres=%#v rollbacks=%d", err, postgresErr, tx.rollbacks)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationPreservesDeferredActivationError(t *testing.T) {
|
|
now := time.Now()
|
|
deferred := &pgconn.PgError{Code: "40001", Message: "could not serialize access"}
|
|
tx := &generationStateTransaction{results: []*jobRows{
|
|
{rows: [][]any{{"ledger-1", int64(1), int64(1), int64(1), int64(1), now, now, int64(-1)}}},
|
|
{err: deferred},
|
|
}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, &generationStatePool{tx: tx})
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job", AmountFen: 1}, json.RawMessage(`{"status":"pending"}`))
|
|
var postgresErr *pgconn.PgError
|
|
if !errors.As(err, &postgresErr) || postgresErr.Code != "40001" || tx.rollbacks != 1 {
|
|
t.Fatalf("err=%v postgres=%#v rollbacks=%d", err, postgresErr, tx.rollbacks)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationTreatsCommittedButLostResponseAsSuccess(t *testing.T) {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
charged := json.RawMessage(`{"status":"charged","amountFen":35,"ledgerEntryId":"ledger-1"}`)
|
|
tx := &generationStateTransaction{
|
|
results: []*jobRows{
|
|
{rows: [][]any{{"ledger-1", int64(965), int64(965), int64(1000), int64(35), now, now, int64(-35)}}},
|
|
{rows: [][]any{{"job-1"}}},
|
|
},
|
|
commitErr: errors.New("connection lost while reading COMMIT response"),
|
|
}
|
|
pool := &generationStatePool{tx: tx, rows: &jobRows{rows: [][]any{{[]byte(charged), now.Format(time.RFC3339Nano)}}}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, pool)
|
|
|
|
got, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job-1", AmountFen: 35}, json.RawMessage(`{"status":"pending","amountFen":35}`))
|
|
if err != nil || string(got) != string(charged) || pool.queries != 1 || tx.rollbacks != 0 {
|
|
t.Fatalf("got=%s err=%v reconciliation queries=%d rollbacks=%d", got, err, pool.queries, tx.rollbacks)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationReturnsDefiniteCommitFailureWhenJobWasNotActivated(t *testing.T) {
|
|
tx := committedGenerationStateTransaction(&pgconn.PgError{Code: "40001", Message: "could not serialize access due to concurrent update"})
|
|
pool := &generationStatePool{tx: tx, rows: &jobRows{}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, pool)
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job-1", AmountFen: 35}, json.RawMessage(`{"status":"pending"}`))
|
|
if err == nil || errors.Is(err, billing.ErrCommitOutcomeUnknown) {
|
|
t.Fatalf("err=%v, want definite ordinary commit failure", err)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationReturnsUnknownOutcomeWhenReconciliationStillShowsPending(t *testing.T) {
|
|
tx := committedGenerationStateTransaction(errors.New("commit response unavailable"))
|
|
pool := &generationStatePool{tx: tx, rows: &jobRows{rows: [][]any{{[]byte(`{"status":"pending","amountFen":35}`), ""}}}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, pool)
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job-1", AmountFen: 35}, json.RawMessage(`{"status":"pending"}`))
|
|
if !errors.Is(err, billing.ErrCommitOutcomeUnknown) {
|
|
t.Fatalf("err=%v, want ErrCommitOutcomeUnknown", err)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationReturnsUnknownOutcomeWhenReconciliationFindsNoJob(t *testing.T) {
|
|
tx := committedGenerationStateTransaction(errors.New("commit response unavailable"))
|
|
pool := &generationStatePool{tx: tx, rows: &jobRows{}}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, pool)
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job-1", AmountFen: 35}, json.RawMessage(`{"status":"pending"}`))
|
|
if !errors.Is(err, billing.ErrCommitOutcomeUnknown) {
|
|
t.Fatalf("err=%v, want ErrCommitOutcomeUnknown", err)
|
|
}
|
|
}
|
|
|
|
func TestChargeAndActivateCreationReturnsUnknownOutcomeWhenReconciliationFails(t *testing.T) {
|
|
tx := committedGenerationStateTransaction(errors.New("commit response unavailable"))
|
|
pool := &generationStatePool{tx: tx, err: errors.New("reconciliation connection unavailable")}
|
|
db := NewDatabase(Config{Backend: BackendPostgres}, pool)
|
|
|
|
_, err := db.ChargeAndActivateCreation(context.Background(), billing.ChargeRequest{OrganizationID: "org", JobID: "job-1", AmountFen: 35}, json.RawMessage(`{"status":"pending"}`))
|
|
if !errors.Is(err, billing.ErrCommitOutcomeUnknown) {
|
|
t.Fatalf("err=%v, want ErrCommitOutcomeUnknown", err)
|
|
}
|
|
}
|
|
|
|
func committedGenerationStateTransaction(commitErr error) *generationStateTransaction {
|
|
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
|
return &generationStateTransaction{
|
|
results: []*jobRows{
|
|
{rows: [][]any{{"ledger-1", int64(965), int64(965), int64(1000), int64(35), now, now, int64(-35)}}},
|
|
{rows: [][]any{{"job-1"}}},
|
|
},
|
|
commitErr: commitErr,
|
|
}
|
|
}
|
|
|
|
type generationStatePool struct {
|
|
tx *generationStateTransaction
|
|
rows *jobRows
|
|
err error
|
|
queries int
|
|
}
|
|
|
|
func (p *generationStatePool) Query(context.Context, string, ...any) (Rows, error) {
|
|
p.queries++
|
|
return p.rows, p.err
|
|
}
|
|
func (p *generationStatePool) Begin(context.Context) (Transaction, error) { return p.tx, nil }
|
|
|
|
type generationStateQuery struct {
|
|
sql string
|
|
args []any
|
|
}
|
|
type generationStateTransaction struct {
|
|
results []*jobRows
|
|
queries []generationStateQuery
|
|
commits, rollbacks int
|
|
commitErr error
|
|
}
|
|
|
|
func (t *generationStateTransaction) Query(_ context.Context, sql string, args ...any) (Rows, error) {
|
|
t.queries = append(t.queries, generationStateQuery{sql: sql, args: append([]any(nil), args...)})
|
|
result := t.results[0]
|
|
t.results = t.results[1:]
|
|
return result, nil
|
|
}
|
|
func (*generationStateTransaction) Exec(context.Context, string, ...any) error { return nil }
|
|
func (t *generationStateTransaction) Commit(context.Context) error { t.commits++; return t.commitErr }
|
|
func (t *generationStateTransaction) Rollback(context.Context) error { t.rollbacks++; return nil }
|