ModelRouter/internal/gateway/stats_test.go

package gateway

import (
	"fmt"
	"os"
	"path/filepath"
	"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))
	}
}