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