zhanlu_proxy/internal/store/store_test.go

package store

import (
	"path/filepath"
	"testing"
	"time"

	"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/stats"
)

func newTestStore(t *testing.T) *Store {
	t.Helper()
	st, err := Open(filepath.Join(t.TempDir(), "test.db"))
	if err != nil {
		t.Fatalf("open store: %v", err)
	}
	t.Cleanup(func() { st.Close() })
	return st
}

func TestCredentialsRoundTrip(t *testing.T) {
	st := newTestStore(t)
	in := auth.Credentials{
		AccessKey: "AK", SecretKey: "SK", Token: "TOK",
		APIKey: "sk-1", ModelBaseURL: "https://up.example", Email: "[email protected]",
		Organization: "org", Team: "team",
	}
	if err := st.SaveCredentials(in); err != nil {
		t.Fatalf("save: %v", err)
	}
	got, err := st.LoadCredentials()
	if err != nil {
		t.Fatalf("load: %v", err)
	}
	if got.APIKey != "sk-1" || got.Email != "[email protected]" || got.Organization != "org" {
		t.Fatalf("round-trip mismatch: %+v", got)
	}
	if got.SavedAt.IsZero() {
		t.Fatalf("saved_at not set")
	}
}

func TestEmptyLoadReturnsZero(t *testing.T) {
	st := newTestStore(t)
	got, err := st.LoadCredentials()
	if err != nil {
		t.Fatalf("load on empty store: %v", err)
	}
	if got.HasAPIKey() {
		t.Fatalf("expected no api key on empty store, got %+v", got)
	}
}

func TestRecordStatsReset(t *testing.T) {
	st := newTestStore(t)
	now := time.Now()
	recs := []stats.Record{
		{Ts: now, Model: "glm-4.7", Stream: false, PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30, CachedTokens: 4, Status: "success", LatencyMs: 5},
		{Ts: now, Model: "glm-4.7", Stream: true, PromptTokens: 5, CompletionTokens: 5, TotalTokens: 10, Status: "success", LatencyMs: 8},
		{Ts: now, Model: "minimax", Stream: false, PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2, Status: "upstream_error", LatencyMs: 3},
	}
	for _, r := range recs {
		if err := st.Record(r); err != nil {
			t.Fatalf("record: %v", err)
		}
	}

	sum, err := st.Stats(stats.Query{})
	if err != nil {
		t.Fatalf("stats: %v", err)
	}
	if sum.Totals.Requests != 3 {
		t.Fatalf("requests = %d, want 3", sum.Totals.Requests)
	}
	if sum.Totals.SuccessRequests != 2 || sum.Totals.ErrorRequests != 1 {
		t.Fatalf("success/error = %d/%d, want 2/1", sum.Totals.SuccessRequests, sum.Totals.ErrorRequests)
	}
	if sum.Totals.TotalTokens != 42 {
		t.Fatalf("total tokens = %d, want 42", sum.Totals.TotalTokens)
	}
	if sum.Totals.CachedTokens != 4 {
		t.Fatalf("cached tokens = %d, want 4", sum.Totals.CachedTokens)
	}
	// prompt total = 10+5+1 = 16, cached = 4 => 0.25
	if sum.Totals.CacheRate < 0.24 || sum.Totals.CacheRate > 0.26 {
		t.Fatalf("cache rate = %v, want ~0.25", sum.Totals.CacheRate)
	}
	if len(sum.PerModel) != 2 {
		t.Fatalf("per-model entries = %d, want 2", len(sum.PerModel))
	}
	// glm-4.7 should lead on total tokens (40 vs 2)
	if sum.PerModel[0].Model != "glm-4.7" || sum.PerModel[0].TotalTokens != 40 || sum.PerModel[0].CachedTokens != 4 {
		t.Fatalf("top model = %+v, want glm-4.7/40/4 cached", sum.PerModel[0])
	}
	if len(sum.Recent) != 3 {
		t.Fatalf("recent entries = %d, want 3", len(sum.Recent))
	}
	// most recent first (id desc) => minimax record
	if sum.Recent[0].Model != "minimax" {
		t.Fatalf("most recent = %+v, want minimax", sum.Recent[0])
	}

	// model filter
	sumF, _ := st.Stats(stats.Query{Model: "minimax"})
	if sumF.Totals.Requests != 1 || sumF.Totals.TotalTokens != 2 {
		t.Fatalf("filtered stats = %+v, want 1/2", sumF.Totals)
	}

	if err := st.Reset(); err != nil {
		t.Fatalf("reset: %v", err)
	}
	sum2, _ := st.Stats(stats.Query{})
	if sum2.Totals.Requests != 0 {
		t.Fatalf("after reset requests = %d, want 0", sum2.Totals.Requests)
	}
	// credentials must survive a stats reset
	creds, _ := st.LoadCredentials()
	_ = creds
}