162 lines
7.0 KiB
Go
162 lines
7.0 KiB
Go
package repository_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"billing/internal/collection"
|
|
"billing/internal/repository"
|
|
)
|
|
|
|
func TestSQLiteUsageQueryPaginatesAndFilters(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
base := time.Date(2026, 8, 15, 0, 0, 0, 0, time.UTC)
|
|
for index := 0; index < 235; index++ {
|
|
failed := index%11 == 0
|
|
record := collection.Record{
|
|
ManagedKeyID: fmt.Sprintf("key-%d", index%3), RequestID: fmt.Sprintf("request-%03d", index),
|
|
ExecutionID: fmt.Sprintf("execution-%03d", index), RequestedAt: base.Add(time.Duration(index) * time.Second),
|
|
Model: fmt.Sprintf("model-%d", index%2), Failed: failed, AuthID: fmt.Sprintf("auth-%d", index%2),
|
|
Endpoint: "/v1/responses", TotalTokens: int64(index + 1),
|
|
}
|
|
if err := store.Insert(context.Background(), record); err != nil {
|
|
t.Fatalf("insert %d: %v", index, err)
|
|
}
|
|
}
|
|
|
|
first, err := store.QueryUsage(context.Background(), collection.UsageQuery{Page: 1, PageSize: 100})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(first.Records) != 100 || first.Total != 235 || first.TotalPages != 3 || first.Page != 1 || first.NextCursor == "" || first.Records[0].RequestID != "request-234" {
|
|
t.Fatalf("unexpected first page: %+v first=%+v", first, first.Records[0])
|
|
}
|
|
second, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: first.NextCursor})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(second.Records) != 100 || second.Page != 2 || second.PreviousCursor == "" || second.Records[0].RequestID != "request-134" {
|
|
t.Fatalf("unexpected second page: %+v first=%+v", second, second.Records[0])
|
|
}
|
|
previous, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: second.PreviousCursor})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(previous.Records) != 100 || previous.Records[0].RequestID != first.Records[0].RequestID || previous.Records[99].RequestID != first.Records[99].RequestID {
|
|
t.Fatalf("previous page did not return the original first page")
|
|
}
|
|
|
|
filters := []struct {
|
|
name string
|
|
query collection.UsageQuery
|
|
check func(collection.Record) bool
|
|
}{
|
|
{name: "key", query: collection.UsageQuery{KeyID: "key-1"}, check: func(record collection.Record) bool { return record.ManagedKeyID == "key-1" }},
|
|
{name: "model", query: collection.UsageQuery{Model: "MODEL-1"}, check: func(record collection.Record) bool { return record.Model == "model-1" }},
|
|
{name: "failed", query: collection.UsageQuery{Result: collection.UsageResultFailed}, check: func(record collection.Record) bool { return record.Failed }},
|
|
{name: "auth", query: collection.UsageQuery{AuthID: "auth-1"}, check: func(record collection.Record) bool { return record.AuthID == "auth-1" }},
|
|
{name: "endpoint", query: collection.UsageQuery{Endpoint: "responses"}, check: func(record collection.Record) bool { return record.Endpoint == "/v1/responses" }},
|
|
{name: "request", query: collection.UsageQuery{RequestID: "request-123"}, check: func(record collection.Record) bool { return record.RequestID == "request-123" }},
|
|
}
|
|
for _, test := range filters {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
test.query.PageSize = 100
|
|
page, err := store.QueryUsage(context.Background(), test.query)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(page.Records) == 0 {
|
|
t.Fatal("filter returned no records")
|
|
}
|
|
for _, record := range page.Records {
|
|
if !test.check(record) {
|
|
t.Fatalf("unexpected filtered record: %+v", record)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
if _, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: first.NextCursor, Model: "model-1"}); err == nil {
|
|
t.Fatal("cursor from another filter unexpectedly accepted")
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageQueryKeepsCursorStableAndProjectsOrphans(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
base := time.Date(2026, 8, 15, 12, 0, 0, 0, time.UTC)
|
|
for index := 0; index < 4; index++ {
|
|
if err := store.Insert(context.Background(), collection.Record{RequestID: fmt.Sprintf("request-%d", index), ExecutionID: fmt.Sprintf("execution-%d", index), RequestedAt: base.Add(time.Duration(index) * time.Second), Model: "model", TotalTokens: 10}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
first, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 2})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.Insert(context.Background(), collection.Record{RequestID: "newest", ExecutionID: "newest", RequestedAt: base.Add(10 * time.Second), Model: "model"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
second, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 2, Cursor: first.NextCursor})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if second.Records[0].RequestID != "request-1" || second.Records[1].RequestID != "request-0" {
|
|
t.Fatalf("cursor shifted after a new insert: %+v", second.Records)
|
|
}
|
|
|
|
startedAt := base.Add(time.Minute)
|
|
if err := store.Insert(context.Background(), collection.Record{RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: "000000", Model: "deepseek-v4-flash", TotalTokens: 99}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.UpsertRequest(context.Background(), collection.RequestRecord{RequestID: "orphan-request", RequestedAt: startedAt, CompletedAt: startedAt.Add(time.Second), Model: "deepseek-v4-flash", Stream: true, Outcome: "canceled", StatusCode: 499, Endpoint: "/v1/responses"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
page, err := store.QueryUsage(context.Background(), collection.UsageQuery{RequestID: "orphan-request", PageSize: 100})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if page.Total != 1 || len(page.Records) != 1 || page.Records[0].TotalTokens != 99 || page.Records[0].Outcome != "canceled" || page.Records[0].StatusCode != 499 {
|
|
t.Fatalf("orphan projection mismatch: %+v", page)
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageDashboardUsesAllIndexedHistory(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
location, _ := time.LoadLocation("Asia/Shanghai")
|
|
today := time.Date(2026, 8, 15, 12, 0, 0, 0, location)
|
|
cost := int64(125_000)
|
|
if err := store.Insert(context.Background(), collection.Record{RequestedAt: today.Add(-time.Hour), Model: "model", InputTokens: 10, OutputTokens: 2, TotalTokens: 12, CostMicros: &cost}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
dashboard, err := store.UsageDashboard(context.Background(), today, 7)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if dashboard.Today.Requests != 1 || dashboard.Today.TotalTokens != 12 || len(dashboard.Users) != 1 || len(dashboard.Days) != 7 || len(dashboard.Users[0].Days) != 7 {
|
|
t.Fatalf("unexpected dashboard: %+v", dashboard)
|
|
}
|
|
if dashboard.Days[6].CostMicros != cost {
|
|
t.Fatalf("today cost = %d, want %d", dashboard.Days[6].CostMicros, cost)
|
|
}
|
|
if dashboard.Users[0].Days[6].CostMicros != cost {
|
|
t.Fatalf("user today cost = %d, want %d", dashboard.Users[0].Days[6].CostMicros, cost)
|
|
}
|
|
}
|