125 lines
4.1 KiB
Go
125 lines
4.1 KiB
Go
package repository_test
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
managedaccess "cpa-ext/internal/access"
|
|
"cpa-ext/internal/collection"
|
|
"cpa-ext/internal/repository"
|
|
)
|
|
|
|
func TestManagedKeysBootstrapLifecycleAndHistoricalUsage(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
if err := store.Insert(context.Background(), collection.Record{
|
|
APIKey: "000000", Model: "gpt-5", RequestedAt: time.Now(), TotalTokens: 10,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
key, err := store.BootstrapManagedKey(context.Background(), "default", "000000")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if key.ID != "key_default" || key.Status != managedaccess.StatusActive || !key.AllModels {
|
|
t.Fatalf("bootstrap key = %+v", key)
|
|
}
|
|
records, err := store.ListRecent(context.Background(), 10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(records) != 1 || records[0].ManagedKeyID != key.ID || records[0].KeyAlias != "default" {
|
|
t.Fatalf("historical records = %+v", records)
|
|
}
|
|
key.Status = managedaccess.StatusDisabled
|
|
key.AllModels = false
|
|
key.Models = []string{"gpt-5"}
|
|
if err := store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.ArchiveManagedKey(context.Background(), key.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
archived, err := store.ManagedKeyByID(context.Background(), key.ID)
|
|
if err != nil || archived.Status != managedaccess.StatusArchived {
|
|
t.Fatalf("archived key = %+v, err=%v", archived, err)
|
|
}
|
|
archived.Status = managedaccess.StatusActive
|
|
if err := store.UpdateManagedKey(context.Background(), archived); err == nil {
|
|
t.Fatal("archived key was restored")
|
|
}
|
|
}
|
|
|
|
func TestManagedKeyStatsUseShanghaiDayBoundary(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
if _, err := store.BootstrapManagedKey(context.Background(), "default", "000000"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cost := int64(25)
|
|
for _, requestedAt := range []time.Time{
|
|
time.Date(2026, 8, 14, 15, 59, 59, 0, time.UTC),
|
|
time.Date(2026, 8, 14, 16, 0, 1, 0, time.UTC),
|
|
} {
|
|
if err := store.Insert(context.Background(), collection.Record{
|
|
ManagedKeyID: "key_default", RequestedAt: requestedAt, TotalTokens: 10, CostMicros: &cost,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
shanghai, _ := time.LoadLocation("Asia/Shanghai")
|
|
today := time.Date(2026, 8, 15, 0, 0, 0, 0, shanghai)
|
|
stats, err := store.KeyStats(context.Background(), "key_default", today)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if stats.Total.Requests != 2 || stats.Today.Requests != 1 || stats.Today.TotalTokens != 10 || stats.Total.CostMicros != 50 {
|
|
t.Fatalf("stats = %+v", stats)
|
|
}
|
|
}
|
|
|
|
func TestManagedKeyStatsCountsOrphanUsageAndAliasLifecycleOnce(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
if _, err := store.BootstrapManagedKey(context.Background(), "default", "000000"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
startedAt := time.Date(2026, 8, 14, 17, 41, 30, 0, time.UTC)
|
|
cost := int64(25)
|
|
if err := store.Insert(context.Background(), collection.Record{
|
|
ManagedKeyID: "key_default", RequestedAt: startedAt.Add(4 * time.Millisecond),
|
|
APIKey: "000000", Model: "deepseek-v4-flash", InputTokens: 90,
|
|
OutputTokens: 7, TotalTokens: 97, CostMicros: &cost,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
|
|
ManagedKeyID: "key_default", RequestID: "request-alias", RequestedAt: startedAt,
|
|
CompletedAt: startedAt.Add(time.Second), Model: "deepseek-source-a",
|
|
Outcome: "succeeded", StatusCode: 200,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
stats, err := store.KeyStats(context.Background(), "key_default", startedAt.Add(-time.Hour))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if stats.Total.Requests != 1 || stats.Total.InputTokens != 90 || stats.Total.OutputTokens != 7 ||
|
|
stats.Total.TotalTokens != 97 || stats.Total.CostMicros != 25 {
|
|
t.Fatalf("stats = %+v", stats)
|
|
}
|
|
}
|