NianAIGC/backend/internal/orchestration/settlement_test.go

package orchestration

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
)

func TestSettlementProcessorSettlesSuccessfulSeedanceActualUsage(t *testing.T) {
	state := &settlementStateStub{}
	ledger := &settlementLedgerStub{posting: billing.WalletPosting{LedgerID: "settlement-1", CreatedAt: time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)}}
	advanced := seedanceSettlementJob(`{"usage":{"completionTokens":20000}}`, `{"status":"charged","amountFen":100,"reservedAmountFen":120,"markupMultiplier":1.2,"parameters":{"resolution":"720p"}}`)
	processor := NewSettlementProcessor(settlementProcessorStub{job: advanced}, ledger, state, func() time.Time { return time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC) })

	got, err := processor.Advance(context.Background(), jobs.Job{ID: "job-1", Provider: "seedance", Status: jobs.StatusRunning})
	if err != nil {
		t.Fatal(err)
	}
	if ledger.calls != 1 || ledger.request.DeltaFen != 11 || ledger.request.JobID != "job-1" || ledger.request.OrganizationID != "org-1" {
		t.Fatalf("settlement request = %#v, calls=%d", ledger.request, ledger.calls)
	}
	if ledger.request.Metadata["completionTokens"] != int64(20000) || ledger.request.Metadata["actualAmountFen"] != int64(111) {
		t.Fatalf("settlement metadata = %#v", ledger.request.Metadata)
	}
	var snapshot map[string]any
	if err := json.Unmarshal(got.Billing, &snapshot); err != nil {
		t.Fatal(err)
	}
	usage := snapshot["providerUsage"].(map[string]any)
	if snapshot["settlementStatus"] != "settled" || snapshot["amountFen"] != float64(111) || snapshot["settlementLedgerEntryId"] != "settlement-1" || usage["completionTokens"] != float64(20000) || state.calls != 1 {
		t.Fatalf("billing snapshot = %#v, writes=%d", snapshot, state.calls)
	}
}

func TestSettlementProcessorMarksMissingProviderUsageEstimated(t *testing.T) {
	state := &settlementStateStub{}
	advanced := seedanceSettlementJob(`{"usage":{}}`, `{"status":"charged","amountFen":100,"markupMultiplier":1.2,"parameters":{"resolution":"720p"}}`)
	processor := NewSettlementProcessor(settlementProcessorStub{job: advanced}, &settlementLedgerStub{}, state, func() time.Time { return time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC) })

	got, err := processor.Advance(context.Background(), jobs.Job{ID: "job-1", Provider: "seedance", Status: jobs.StatusRunning})
	if err != nil {
		t.Fatal(err)
	}
	var snapshot map[string]any
	_ = json.Unmarshal(got.Billing, &snapshot)
	if snapshot["settlementStatus"] != "estimated" || snapshot["settlementReason"] != "provider_usage_unavailable" || state.calls != 1 {
		t.Fatalf("billing snapshot = %#v, writes=%d", snapshot, state.calls)
	}
}

func TestSettlementProcessorQuotaExemptUpdatesSnapshotWithoutWalletAndIsIdempotent(t *testing.T) {
	state := &settlementStateStub{}
	ledger := &settlementLedgerStub{}
	advanced := seedanceSettlementJob(`{"usage":{"completionTokens":20000}}`, `{"status":"not_charged","amountFen":100,"markupMultiplier":1.2,"quotaExempt":true,"parameters":{"resolution":"720p"}}`)
	processor := NewSettlementProcessor(settlementProcessorStub{job: advanced}, ledger, state, time.Now)

	got, err := processor.Advance(context.Background(), jobs.Job{ID: "job-1", Provider: "seedance", Status: jobs.StatusRunning})
	if err != nil || ledger.calls != 0 || state.calls != 1 {
		t.Fatalf("first settlement calls ledger=%d state=%d err=%v", ledger.calls, state.calls, err)
	}
	processor.inner = settlementProcessorStub{job: got}
	got, err = processor.Advance(context.Background(), got)
	if err != nil || ledger.calls != 0 || state.calls != 1 {
		t.Fatalf("repeat settlement calls ledger=%d state=%d err=%v", ledger.calls, state.calls, err)
	}
	var snapshot map[string]any
	_ = json.Unmarshal(got.Billing, &snapshot)
	if snapshot["amountFen"] != float64(111) || snapshot["settlementStatus"] != "settled" {
		t.Fatalf("billing snapshot = %#v", snapshot)
	}
}

func TestSettlementProcessorLeavesNonSeedanceSuccessAlone(t *testing.T) {
	state := &settlementStateStub{}
	advanced := jobs.Job{ID: "job-2", Provider: "bailian", Status: jobs.StatusSucceeded, Billing: json.RawMessage(`{"status":"charged"}`)}
	got, err := NewSettlementProcessor(settlementProcessorStub{job: advanced}, &settlementLedgerStub{}, state, time.Now).Advance(context.Background(), jobs.Job{ID: "job-2"})
	if err != nil || got.Provider != "bailian" || state.calls != 0 {
		t.Fatalf("job=%#v writes=%d err=%v", got, state.calls, err)
	}
}

func seedanceSettlementJob(response, billingJSON string) jobs.Job {
	return jobs.Job{
		ID: "job-1", Provider: "seedance", Capability: "video.generate", Status: jobs.StatusSucceeded,
		RequestPayload:  json.RawMessage(`{"settings":{"resolution":"720p"},"inputUrls":[]}`),
		ResponsePayload: json.RawMessage(response), Billing: json.RawMessage(billingJSON),
		UsageContext: json.RawMessage(`{"organizationId":"org-1","accountId":"account-1"}`),
	}
}

type settlementProcessorStub struct {
	job jobs.Job
	err error
}

func (s settlementProcessorStub) Advance(context.Context, jobs.Job) (jobs.Job, error) {
	return s.job, s.err
}

type settlementLedgerStub struct {
	request billing.SettlementRequest
	posting billing.WalletPosting
	err     error
	calls   int
}

func (s *settlementLedgerStub) Settle(_ context.Context, request billing.SettlementRequest) (billing.WalletPosting, error) {
	s.calls++
	s.request = request
	return s.posting, s.err
}

type settlementStateStub struct {
	id      string
	billing json.RawMessage
	err     error
	calls   int
}

func (s *settlementStateStub) WriteBilling(_ context.Context, id string, value json.RawMessage) error {
	s.calls++
	s.id = id
	s.billing = append(json.RawMessage(nil), value...)
	return s.err
}