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