倍率问题和图片问题调整

backend/internal/postgres/billing_defaults_test.go

@@ -3,11 +3,16 @@ 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) {
@@ -31,3 +36,123 @@ func TestSeedBillingPriceRulesUsesOneIdempotentUpsert(t *testing.T) {
 		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...)
+}