ModelRouter/internal/gateway/stats_test.go

package gateway

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestStatsByStatus(t *testing.T) {
	s := NewStats(100)
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 402, OK: false})
	s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 400, OK: false})
	snap := s.Snapshot(0, "")
	bs, ok := snap["by_status"].([]agrRow)
	if !ok {
		t.Fatalf("by_status missing: %#v", snap["by_status"])
	}
	if len(bs) != 3 {
		t.Fatalf("want 3 status buckets, got %d: %#v", len(bs), bs)
	}
	if bs[0].Name != "200" || bs[0].OK != 1 || bs[0].Err != 0 {
		t.Fatalf("bucket 200 wrong: %#v", bs[0])
	}
	if bs[1].Name != "400" || bs[1].Err != 1 {
		t.Fatalf("bucket 400 wrong: %#v", bs[1])
	}
	if bs[2].Name != "402" || bs[2].Err != 1 {
		t.Fatalf("bucket 402 wrong: %#v", bs[2])
	}
}

func TestAuditRotation(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	s := NewStats(10)
	s.LoadAudit(path)

	oldRotate, oldKeep := auditRotateBytes, auditKeepOld
	auditRotateBytes, auditKeepOld = 64, 10
	defer func() { auditRotateBytes, auditKeepOld = oldRotate, oldKeep }()

	oldFiles := func() []string {
		matches, _ := filepath.Glob(path + ".*.old")
		return matches
	}

	for i := 0; i < 3; i++ {
		s.AppendAudit("ev", map[string]interface{}{"i": i})
	}
	if got := len(oldFiles()); got != 1 {
		t.Fatalf("want 1 rotated file after first overflow, got %d", got)
	}
	if b, err := os.ReadFile(path); err != nil || len(b) == 0 {
		t.Fatalf("active audit file must continue appending: %v %d bytes", err, len(b))
	}

	// seed 12 fake old files; the next rotation must prune back to keep=10
	for i := 1; i <= 12; i++ {
		name := fmt.Sprintf("%s.%010d.old", path, i)
		_ = os.WriteFile(name, []byte("x\n"), 0644)
	}
	s.AppendAudit("ev", map[string]interface{}{"i": 98})
	s.AppendAudit("ev", map[string]interface{}{"i": 99})
	if got := len(oldFiles()); got != auditKeepOld {
		t.Fatalf("want keeper %d old files, got %d", auditKeepOld, got)
	}
}

func TestAuditRotationRecords(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	s := NewStats(10)
	s.LoadAudit(path)

	oldRotate := auditRotateBytes
	auditRotateBytes = 64
	defer func() { auditRotateBytes = oldRotate }()

	for i := 0; i < 5; i++ {
		s.Record(Req{Key: "k", Model: "m", Source: "s", Status: 200, OK: true})
	}
	matches, _ := filepath.Glob(path + ".*.old")
	if len(matches) != 1 {
		t.Fatalf("Record must rotate too: got %d old files", len(matches))
	}
}

func TestLoadAuditFullReplay(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	lines := []string{
		`{"obj":"access","time":1699999999000,"key":"k","method":"GET","path":"/api/stats","status":200}`,
		`{"time":1700000000000,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`,
		`{this is not valid json`,
		`{"time":1700003600000,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`,
		"garbage-not-json\n",
		`{"time":1700007200000,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`,
	}
	loaded := strings.Join(lines, "\n") + "\n" + strings.Repeat("x", 1<<18) + "\n"
	// oversized row at the END proves the scanner tolerates >64KB lines and
	// still finishes the replay instead of truncating silently.
	loaded += `{"time":1700010800000,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}` + "\n"
	if err := os.WriteFile(path, []byte(loaded), 0644); err != nil {
		t.Fatal(err)
	}

	s := NewStats(1000)
	s.LoadAudit(path)

	got := s.byModel["m"]
	if got == nil || got.Tokens != 100+50+200+20+1+1 {
		t.Fatalf("aggregates must replay EVERY request row, got %#v", got)
	}
	if s.byModel["m2"] == nil || s.byModel["m2"].Tokens != 10 {
		t.Fatalf("m2 must be replayed too, got %#v", s.byModel["m2"])
	}
	// access rows are not requests: 4 real rows, junk skipped
	if len(s.recs) != 4 {
		t.Fatalf("ring must hold only real requests, got %d rows: %#v", len(s.recs), s.recs)
	}
	// quota window rebuilt from full history
	if w := s.WindowTokens("m", "s", hourSec); w != 372 {
		t.Fatalf("window tokens want 372, got %d", w)
	}
	// by_status only from requests (200 x3, 503 x1) — access line must not count
	if st := s.byStatus[200]; st == nil || st.Reqs != 3 {
		t.Fatalf("by_status 200 want 3 reqs, got %#v", st)
	}
	if s.byStatus[200].Err != 0 || s.byStatus[503] == nil || s.byStatus[503].Reqs != 1 {
		t.Fatalf("by_status wrong: %#v", s.byStatus)
	}
}