openteam/server/internal/usage/usage_test.go

package usage

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

	"github.com/openteam/server/internal/store"
	"github.com/shopspring/decimal"
	"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)
	cost := decimal.NewFromFloat(0.15).Add(decimal.NewFromFloat(0.60)).
		Mul(decimal.NewFromInt(1000)).Div(decimal.NewFromInt(1e6)).Round(8)
	rec.Record(&store.UsageLog{
		TraceID: "bill", UserID: u, ModelID: m, ModelName: "m1",
		InputTokens: 1000, OutputTokens: 1000,
		InputPrice: decimal.NewFromFloat(0.15), OutputPrice: decimal.NewFromFloat(0.60),
		Cost:      cost,
		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 := decimal.NewFromInt(1).Sub(cost)
	if !user.Balance.Equal(want) {
		t.Fatalf("余额应精确扣到 %s，实际 %s", 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.IsNegative() {
		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: decimal.NewFromFloat(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: decimal.NewFromFloat(0.15), OutputPrice: decimal.NewFromFloat(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: decimal.NewFromInt(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.Equal(decimal.NewFromInt(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: decimal.NewFromFloat(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 := decimal.NewFromInt(100).Sub(decimal.NewFromFloat(0.01).Mul(decimal.NewFromInt(n))); !user.Balance.Equal(want) {
		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 want := decimal.NewFromFloat(0.01).Mul(decimal.NewFromInt(n)); !daily[0].Cost.Equal(want) {
		t.Fatalf("日聚合 cost = %v，期望 %v", daily[0].Cost, want)
	}
	// 流水里的 balance_after 应与扣款后余额一致（批内累计）
	var last store.BalanceLog
	db.Where("user_id = ?", u).Order("id DESC").First(&last)
	if !last.BalanceAfter.Equal(user.Balance) {
		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: decimal.NewFromInt(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: decimal.NewFromFloat(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.GreaterThan(logs[i-1].BalanceAfter) {
				t.Fatalf("用户 %d 的流水余额应单调递减（顺序错乱）：%v → %v",
					u, logs[i-1].BalanceAfter, logs[i].BalanceAfter)
			}
		}
	}
}

// 渠道×小时桶：状态计数含 error 行，token/成本只累计可计费行，跨小时分桶。
func TestChannelHourlyBucketAggregated(t *testing.T) {
	db := openTestDB(t)
	u := seedUser(t, db, 100)
	m := seedModel(t, db)

	rec := NewRecorder(db)
	now := time.Now().UTC()
	hour := now.Format("2006-01-02T15")
	prevHour := now.Add(-time.Hour).Format("2006-01-02T15")
	mk := func(trace, status string, at time.Time, ch uint64) *store.UsageLog {
		return &store.UsageLog{TraceID: trace, UserID: u, ModelID: m, ModelName: "m1",
			ChannelID: ch, InputTokens: 100, OutputTokens: 50, Cost: decimal.NewFromFloat(0.02),
			Status: status, LatencyMS: 200, CreatedAt: at}
	}
	rec.Record(mk("h1", store.UsageStatusSuccess, now, 7))
	rec.Record(mk("h2", store.UsageStatusSuccess, now, 7))
	rec.Record(mk("h3", store.UsageStatusError, now, 7))
	rec.Record(mk("h4", store.UsageStatusSuccess, now.Add(-time.Hour), 7))
	rec.Record(mk("h5", store.UsageStatusSuccess, now, 0)) // channel_id=0 不入桶
	rec.Close()

	var buckets []store.UsageChannelHourly
	db.Where("channel_id = ?", 7).Order("hour DESC").Find(&buckets)
	if len(buckets) != 2 {
		t.Fatalf("应有两个小时桶，实际 %d", len(buckets))
	}
	cur, prev := buckets[0], buckets[1]
	if cur.Hour != hour || prev.Hour != prevHour {
		t.Fatalf("桶小时错误：%q / %q，期望 %q / %q", cur.Hour, prev.Hour, hour, prevHour)
	}
	if cur.Requests != 3 || cur.Success != 2 || cur.Errors != 1 {
		t.Fatalf("当前桶计数 = %d/%d/%d，期望 3/2/1（requests/success/errors）", cur.Requests, cur.Success, cur.Errors)
	}
	// 只有 success 行计延迟
	if cur.LatencySum != 400 || cur.LatencyMax != 200 {
		t.Fatalf("延迟聚合 = sum %d max %d，期望 400/200", cur.LatencySum, cur.LatencyMax)
	}
	// token/成本只累计可计费行（2×success）
	if cur.InputTokens != 200 || cur.Cost.String() != "0.04" {
		t.Fatalf("token/成本 = %d/%s，期望 200/0.04", cur.InputTokens, cur.Cost)
	}
	if prev.Requests != 1 || prev.Success != 1 {
		t.Fatalf("上一小时桶计数 = %d/%d，期望 1/1", prev.Requests, prev.Success)
	}
}

// 分片函数：同一用户恒定同片，且覆盖所有分片。
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)
	}
}