openteam/server/internal/usage/usage_test.go

package usage

import (
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/openteam/server/internal/store"
	"gorm.io/gorm"
)

func TestCloseDrainsQueue(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 100)

	rec := NewRecorder(db)
	// 快速灌入超过一批的记录（batchSize=32），大部分会留在 channel 里
	for i := 0; i < 100; i++ {
		rec.Record(&store.UsageLog{
			TraceID: fmt.Sprintf("t%d", i), UserID: u, ModelName: "m1",
			Status: store.UsageStatusError, ErrorCode: strPtr("test"), // error 状态不触发扣费分支
			CreatedAt: time.Now().UTC(),
		})
	}
	rec.Close()

	var n int64
	db.Table("usage_logs").Count(&n)
	if n != 100 {
		t.Fatalf("Close 应排空队列全部落库，实际 %d/100", n)
	}
}

func TestRecordAfterCloseDoesNotPanic(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 100)

	rec := NewRecorder(db)
	rec.Close()
	// Shutdown 超时后被强断的在途请求仍可能调 Record：不允许 send-on-closed panic
	done := make(chan struct{})
	go func() {
		defer close(done)
		rec.Record(&store.UsageLog{
			TraceID: "late", UserID: u, ModelName: "m1",
			Status: store.UsageStatusError, ErrorCode: strPtr("late"),
			CreatedAt: time.Now().UTC(),
		})
	}()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Record 在 Close 之后调用发生阻塞（疑似 panic 被吞）")
	}

	var n int64
	db.Table("usage_logs").Where("trace_id = ?", "late").Count(&n)
	if n != 1 {
		t.Fatalf("Close 后的 Record 应同步落库，实际 %d 条", n)
	}
}

func TestCloseIsIdempotent(t *testing.T) {
	db := openTestDB(t)
	seedUser(t, db, 0)
	rec := NewRecorder(db)
	rec.Close()
	rec.Close() // serve.go 的 defer Shutdown 与显式 Close 双调用
}

func TestBillingDeductsBalanceAtomically(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 1)
	m := seedModel(t, db)

	rec := NewRecorder(db)
	rec.Record(&store.UsageLog{
		TraceID: "bill", UserID: u, ModelID: m, ModelName: "m1",
		InputTokens: 1000, OutputTokens: 1000,
		InputPrice: 0.15, OutputPrice: 0.60,
		Cost:      0.15*1000/1e6 + 0.60*1000/1e6,
		Status:    store.UsageStatusSuccess,
		LatencyMS: 5,
		CreatedAt: time.Now().UTC(),
	})
	rec.Close()

	var user store.User
	if err := db.First(&user, u).Error; err != nil {
		t.Fatal(err)
	}
	want := 1 - (0.15*1000/1e6 + 0.60*1000/1e6)
	if diff := user.Balance - want; diff > 1e-9 || diff < -1e-9 {
		t.Fatalf("余额应扣到 %.10f，实际 %.10f", want, user.Balance)
	}
	var bl store.BalanceLog
	if err := db.Where("user_id = ?", u).First(&bl).Error; err != nil {
		t.Fatalf("余额流水应落库: %v", err)
	}
	if bl.Change >= 0 {
		t.Fatalf("流水应为负数扣费，实际 %v", bl.Change)
	}
}

// ---------------------------------------------------------------------------
// helpers

func openTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	dsn := "/tmp/usage_test_" + t.Name() + "_" + time.Now().Format("150405.000000000") + ".db"
	db, err := store.Open("sqlite", dsn)
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	t.Cleanup(func() {
		sqlDB, _ := db.DB()
		sqlDB.Close()
		os.Remove(dsn)
	})
	return db
}

func seedUser(t *testing.T, db *gorm.DB, balance float64) uint64 {
	t.Helper()
	u := store.User{Username: "u" + fmt.Sprint(time.Now().UnixNano()), Email: "u@x.com",
		PasswordHash: "x", Role: store.RoleUser, Balance: balance, Status: store.UserStatusActive}
	if err := db.Create(&u).Error; err != nil {
		t.Fatal(err)
	}
	return u.ID
}

func seedModel(t *testing.T, db *gorm.DB) uint64 {
	t.Helper()
	m := store.Model{Name: "m1", InputPrice: 0.15, OutputPrice: 0.60, Enabled: true}
	if err := db.Create(&m).Error; err != nil {
		t.Fatal(err)
	}
	return m.ID
}

func strPtr(s string) *string { return &s }

// 同一 trace 重复记账（重试/重复 flush 场景）只扣一次款：流水 ref_id 有唯一索引，
// 但扣款本身没有，billUser 必须先查再扣。
func TestBillingRefIDImpotency(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 5)

	rec := NewRecorder(db)
	mk := func(trace string) *store.UsageLog {
		return &store.UsageLog{TraceID: trace, UserID: u, ModelName: "m1",
			Status: store.UsageStatusSuccess, Cost: 1, CreatedAt: time.Now().UTC()}
	}
	rec.Record(mk("dup"))
	rec.Record(mk("dup"))
	rec.Close()

	var user store.User
	if err := db.First(&user, u).Error; err != nil {
		t.Fatal(err)
	}
	if user.Balance != 4 {
		t.Errorf("balance = %v, want 4 (重复 trace 不重复扣款)", user.Balance)
	}
	var n int64
	db.Model(&store.BalanceLog{}).Where("user_id = ?", u).Count(&n)
	if n != 1 {
		t.Errorf("balance logs = %d, want 1", n)
	}
}

// 一批同用户请求：一次扣款、一条日聚合（而不是逐行 upsert）。
func TestBillingAggregatesBatch(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 100)
	m := seedModel(t, db)

	rec := NewRecorder(db)
	const n = 32
	for i := 0; i < n; i++ {
		rec.Record(&store.UsageLog{
			TraceID: fmt.Sprintf("agg-%d", i), UserID: u, ModelID: m, ModelName: "m1",
			InputTokens: 10, OutputTokens: 5, Cost: 0.01,
			Status: store.UsageStatusSuccess, CreatedAt: time.Now().UTC(),
		})
	}
	rec.Close()

	var user store.User
	if err := db.First(&user, u).Error; err != nil {
		t.Fatal(err)
	}
	if want := 100 - 0.01*n; user.Balance-want > 1e-9 || user.Balance-want < -1e-9 {
		t.Fatalf("余额 = %v，期望 %v", user.Balance, want)
	}
	var ledgers int64
	db.Model(&store.BalanceLog{}).Where("user_id = ?", u).Count(&ledgers)
	if ledgers != n {
		t.Fatalf("流水条数 = %d，期望 %d（每请求一条明细）", ledgers, n)
	}
	var daily []store.UsageDaily
	db.Where("user_id = ?", u).Find(&daily)
	if len(daily) != 1 {
		t.Fatalf("日聚合应只有 1 行，实际 %d", len(daily))
	}
	if daily[0].Requests != n {
		t.Fatalf("日聚合 requests = %d，期望 %d", daily[0].Requests, n)
	}
	if d := daily[0].Cost - 0.01*n; d > 1e-9 || d < -1e-9 {
		t.Fatalf("日聚合 cost = %v，期望 %v", daily[0].Cost, 0.01*n)
	}
	// 流水里的 balance_after 应与扣款后余额一致（批内累计）
	var last store.BalanceLog
	db.Where("user_id = ?", u).Order("id DESC").First(&last)
	if last.BalanceAfter-user.Balance > 1e-9 || user.Balance-last.BalanceAfter > 1e-9 {
		t.Fatalf("最后一条流水 balance_after = %v，实际余额 %v", last.BalanceAfter, user.Balance)
	}
}

// 多 worker 落库：同一用户恒定落在同一分片（流水顺序与提交顺序一致），
// 不同用户并行落库，最终一条不丢。
func TestRecorderShardedByUser(t *testing.T) {
	db := openTestDB(t)
	var users []uint64
	for i := 0; i < 3; i++ {
		u := store.User{Username: fmt.Sprintf("shard-u%d", i), Email: fmt.Sprintf("shard%d@x.com", i),
			PasswordHash: "x", Role: store.RoleUser, Balance: 100, Status: store.UserStatusActive}
		if err := db.Create(&u).Error; err != nil {
			t.Fatal(err)
		}
		users = append(users, u.ID)
	}
	// 显式用 3 个分片：SQLite 生产配置是单 worker，但分片语义本身仍要验证
	rec := newRecorder(db, 3)
	const per = 40
	for i := 0; i < per; i++ {
		for _, u := range users {
			rec.Record(&store.UsageLog{
				TraceID: fmt.Sprintf("u%d-%d", u, i), UserID: u, ModelName: "m1",
				Cost: 0.001, Status: store.UsageStatusSuccess, CreatedAt: time.Now().UTC(),
			})
		}
	}
	rec.Close()

	var n int64
	db.Model(&store.UsageLog{}).Count(&n)
	if n != int64(per*len(users)) {
		t.Fatalf("落库 %d 条，期望 %d（一条不丢）", n, per*len(users))
	}
	// 每个用户的分片一致 → 流水按 id 递增的顺序应与 trace 序号一致
	for _, u := range users {
		var logs []store.BalanceLog
		db.Where("user_id = ?", u).Order("id ASC").Find(&logs)
		if len(logs) != per {
			t.Fatalf("用户 %d 流水 %d 条，期望 %d", u, len(logs), per)
		}
		for i := 1; i < len(logs); i++ {
			if logs[i].BalanceAfter > logs[i-1].BalanceAfter {
				t.Fatalf("用户 %d 的流水余额应单调递减（顺序错乱）：%v → %v",
					u, logs[i-1].BalanceAfter, logs[i].BalanceAfter)
			}
		}
	}
}

// 分片函数：同一用户恒定同片，且覆盖所有分片。
func TestShardOfStable(t *testing.T) {
	rec := NewRecorder(openTestDB(t))
	defer rec.Close()
	for _, u := range []uint64{1, 2, 3, 17, 1 << 40} {
		if rec.shardOf(u) != rec.shardOf(u) {
			t.Fatalf("用户 %d 的分片不稳定", u)
		}
	}
	seen := map[chan *store.UsageLog]bool{}
	for u := uint64(0); u < 64; u++ {
		seen[rec.shardOf(u)] = true
	}
	if len(seen) != rec.workers {
		t.Fatalf("只用到 %d 个分片，期望 %d", len(seen), rec.workers)
	}
}