NianAIGC/backend/internal/postgres/billing_defaults_test.go

package postgres

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
	"github.com/jackc/pgx/v5"
)

func TestSeedBillingPriceRulesUsesOneIdempotentUpsert(t *testing.T) {
	q := &fakeQuerier{}
	db := NewDatabase(Config{Backend: BackendPostgres}, q)
	rules := []billing.PriceRule{{ID: "r", Provider: "seedance", Capability: "video.generate", ReqKey: "m", VariantKey: "resolution=720p", Unit: billing.UnitVideoSecond, StandardUnitPriceFen: 99, MarkupMultiplier: 1.2, Enabled: true}}
	if err := db.SeedBillingPriceRules(context.Background(), rules); err != nil {
		t.Fatal(err)
	}
	if q.sql != SeedBillingPriceRulesSQL || !strings.Contains(q.sql, "ON CONFLICT (provider, capability, COALESCE(req_key, ''), COALESCE(variant_key, ''), COALESCE(conditions, '{}'::jsonb)) DO UPDATE") {
		t.Fatalf("sql = %q", q.sql)
	}
	if len(q.args) != 1 {
		t.Fatalf("args = %#v", q.args)
	}
	var decoded []billing.PriceRule
	if err := json.Unmarshal(q.args[0].([]byte), &decoded); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(decoded, rules) {
		t.Fatalf("payload = %#v", decoded)
	}
}

// Runs against a private temporary PostgreSQL cluster when the local binaries
// are supplied. No application database or credentials are used.
func TestSeedBillingPriceRulesPreservesTierMarkupInPostgres(t *testing.T) {
	binDir := os.Getenv("NIAN_TEST_POSTGRES_BIN_DIR")
	if binDir == "" {
		t.Skip("set NIAN_TEST_POSTGRES_BIN_DIR to run the PostgreSQL integration test")
	}
	base, err := os.MkdirTemp("/tmp", "nian-billing-pg-")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = os.RemoveAll(base) })
	dataDir := filepath.Join(base, "data")
	runPG := func(name string, args ...string) {
		t.Helper()
		command := exec.Command(filepath.Join(binDir, name), args...)
		if output, err := command.CombinedOutput(); err != nil {
			t.Fatalf("%s: %v: %s", name, err, output)
		}
	}
	runPG("initdb", "-D", dataDir, "-A", "trust", "-U", "postgres")
	options := fmt.Sprintf("-F -c unix_socket_directories=%s -c listen_addresses='' -p 55432", base)
	runPG("pg_ctl", "-D", dataDir, "-l", filepath.Join(base, "server.log"), "-o", options, "-w", "start")
	t.Cleanup(func() {
		command := exec.Command(filepath.Join(binDir, "pg_ctl"), "-D", dataDir, "-m", "immediate", "-w", "stop")
		if output, err := command.CombinedOutput(); err != nil {
			t.Errorf("stop temporary PostgreSQL: %v: %s", err, output)
		}
	})
	ctx := context.Background()
	conn, err := pgx.Connect(ctx, fmt.Sprintf("host=%s port=55432 user=postgres dbname=postgres sslmode=disable", base))
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close(ctx)
	_, err = conn.Exec(ctx, `CREATE TABLE public.billing_price_rules (
		id text PRIMARY KEY, provider text NOT NULL, capability text NOT NULL, req_key text,
		variant_key text, unit text NOT NULL, standard_unit_price_fen bigint NOT NULL,
		markup_multiplier numeric NOT NULL, enabled boolean NOT NULL, conditions jsonb NOT NULL,
		quantity_source text, priority integer NOT NULL, note text, source jsonb,
		parameter_dimensions jsonb NOT NULL, updated_at timestamptz DEFAULT now());
		CREATE UNIQUE INDEX billing_price_rules_match_idx ON public.billing_price_rules
		(provider, capability, coalesce(req_key, ''), coalesce(variant_key, ''), coalesce(conditions, '{}'::jsonb));`)
	if err != nil {
		t.Fatal(err)
	}
	rules := make([]billing.PriceRule, 0, 2)
	for _, rule := range billing.DefaultBillingPriceRules() {
		if rule.ID == "base-seedream-5-0-pro" || rule.ID == "base-seedream-5-0-pro-layers" {
			rules = append(rules, rule)
		}
	}
	if len(rules) != 2 {
		t.Fatalf("Seedream defaults: %d rules", len(rules))
	}
	db := NewDatabase(Config{Backend: BackendPostgres}, seedIntegrationQuerier{conn})
	if err := db.SeedBillingPriceRules(ctx, rules); err != nil {
		t.Fatal(err)
	}
	_, err = conn.Exec(ctx, `UPDATE public.billing_price_rules
	SET parameter_dimensions = jsonb_set(jsonb_set(parameter_dimensions,
		'{0,tiers,1,markupMultiplier}', '"bad"'::jsonb), '{0,tiers,2,markupMultiplier}', '1.3'::jsonb),
		markup_multiplier = 1.5
	WHERE id = 'base-seedream-5-0-pro'`)
	if err != nil {
		t.Fatal(err)
	}
	_, err = conn.Exec(ctx, `UPDATE public.billing_price_rules
	SET parameter_dimensions = jsonb_set(parameter_dimensions, '{0,tiers,2,markupMultiplier}', '1.4'::jsonb)
	WHERE id = 'base-seedream-5-0-pro-layers'`)
	if err != nil {
		t.Fatal(err)
	}
	for i := range rules {
		rules[i].StandardUnitPriceFen++
		rules[i].Dimensions[0].Tiers = append(rules[i].Dimensions[0].Tiers[1:], billing.ParameterTier{
			Value: "3K", Label: "3K", StandardFactor: 3, MarkupMultiplier: 1.2, Enabled: true,
		})
		rules[i].Dimensions[0].Tiers[1].StandardFactor = 1.5
	}
	for attempt := 0; attempt < 2; attempt++ {
		if err := db.SeedBillingPriceRules(ctx, rules); err != nil {
			t.Fatal(err)
		}
	}
	for _, rule := range rules {
		var cost int64
		var topLevelMarkup float64
		var dimensionsJSON []byte
		if err := conn.QueryRow(ctx, `SELECT standard_unit_price_fen, markup_multiplier, parameter_dimensions FROM public.billing_price_rules WHERE id = $1`, rule.ID).Scan(&cost, &topLevelMarkup, &dimensionsJSON); err != nil {
			t.Fatal(err)
		}
		var dimensions []billing.ParameterDimension
		if err := json.Unmarshal(dimensionsJSON, &dimensions); err != nil {
			t.Fatal(err)
		}
		if cost != rule.StandardUnitPriceFen || len(dimensions) != 1 || len(dimensions[0].Tiers) != 4 {
			t.Fatalf("%s: cost=%d dimensions=%+v", rule.ID, cost, dimensions)
		}
		if dimensions[0].Tiers[1].StandardFactor != 1.5 || dimensions[0].Tiers[3].Value != "3K" || dimensions[0].Tiers[3].MarkupMultiplier != 1.2 {
			t.Fatalf("%s: default tier definitions not refreshed: %+v", rule.ID, dimensions[0].Tiers)
		}
		wantMarkup := 1.3
		if rule.ID == "base-seedream-5-0-pro-layers" {
			wantMarkup = 1.4
		} else if topLevelMarkup != 1.5 || dimensions[0].Tiers[0].MarkupMultiplier != 1.2 {
			t.Fatalf("basic rule lost top-level markup or malformed tier did not fall back: %+v", dimensions[0].Tiers)
		}
		if dimensions[0].Tiers[1].MarkupMultiplier != wantMarkup {
			t.Fatalf("%s: tier markup = %g, want %g", rule.ID, dimensions[0].Tiers[1].MarkupMultiplier, wantMarkup)
		}
	}
}

type seedIntegrationQuerier struct{ conn *pgx.Conn }

func (q seedIntegrationQuerier) Query(ctx context.Context, sql string, args ...any) (Rows, error) {
	return q.conn.Query(ctx, sql, args...)
}