ModelRouter/internal/gateway/stats_period_test.go

package gateway

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// auditPeriod writes records straight into an audit file, which is what the
// period aggregator reads. Going through Record() would also rotate/prune the
// file; here the file shape is the thing under test.
func auditPeriod(t *testing.T, path string, recs ...Req) {
	t.Helper()
	f, err := os.Create(path)
	if err != nil {
		t.Fatalf("create audit file: %v", err)
	}
	defer f.Close()
	for _, r := range recs {
		if r.Time == 0 {
			t.Fatal("test record needs an explicit timestamp")
		}
		if r.Type == "" {
			r.Type = "chat"
		}
		b, err := json.Marshal(r)
		if err != nil {
			t.Fatalf("marshal record: %v", err)
		}
		if _, err := f.Write(append(b, '\n')); err != nil {
			t.Fatalf("write record: %v", err)
		}
	}
}

func periodStats(t *testing.T, auditPath string) *Stats {
	t.Helper()
	s := NewStats(10)
	s.mu.Lock()
	s.auditPath = auditPath
	s.mu.Unlock()
	return s
}

// TestPeriodWindowBoundaries is the boundary contract: a day window covers
// exactly today-from-midnight, a week window starts Monday (not Sunday — the
// Go Weekday() trap), and a month window starts on the 1st. All UTC.
func TestPeriodWindowBoundaries(t *testing.T) {
	// Wednesday 2026-03-11 15:30 UTC.
	now := time.Date(2026, 3, 11, 15, 30, 0, 0, time.UTC)
	day := time.Date(2026, 3, 11, 0, 0, 0, 0, time.UTC)
	if got := time.UnixMilli(periodStart(PeriodDay, now)).UTC(); !got.Equal(day) {
		t.Errorf("day window starts %s, want %s", got, day)
	}
	// Wednesday minus 2 days = Monday the 9th.
	monday := time.Date(2026, 3, 9, 0, 0, 0, 0, time.UTC)
	if got := time.UnixMilli(periodStart(PeriodWeek, now)).UTC(); !got.Equal(monday) {
		t.Errorf("week window starts %s, want Monday %s", got, monday)
	}
	// Same instant on a Sunday must still start on the PRECEDING Monday,
	// never on the Sunday itself. This is the off-by-one that
	// int(now.Weekday()) would introduce (Sunday=0 → no shift).
	sun := time.Date(2026, 3, 15, 10, 0, 0, 0, time.UTC) // Sunday
	prevMon := time.Date(2026, 3, 9, 0, 0, 0, 0, time.UTC)
	if got := time.UnixMilli(periodStart(PeriodWeek, sun)).UTC(); !got.Equal(prevMon) {
		t.Errorf("Sunday week window starts %s, want the Monday before it %s", got, prevMon)
	}
	first := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC)
	if got := time.UnixMilli(periodStart(PeriodMonth, now)).UTC(); !got.Equal(first) {
		t.Errorf("month window starts %s, want %s", got, first)
	}
	if got := periodStart(PeriodAll, now); got != 0 {
		t.Errorf("all window from = %d, want 0 (unbounded)", got)
	}
}

// TestPeriodAggregationExcludesOlderRecords is the reason the aggregator reads
// the audit file: a record outside the window must not be counted, even though
// it sits in the very same file, right next to records that are.
func TestPeriodAggregationExcludesOlderRecords(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	now := time.Date(2026, 3, 11, 15, 0, 0, 0, time.UTC)
	h := func(n int) int64 { return now.Add(-time.Duration(n) * time.Hour).UnixMilli() }

	auditPeriod(t, path,
		// inside today
		Req{Time: h(1), Model: "m", Source: "s", Prompt: 100, Compl: 10, LatMs: 5, OK: true},
		Req{Time: h(2), Model: "m", Source: "s", Prompt: 200, Compl: 20, LatMs: 7, OK: true},
		// outside today (yesterday) but inside the week
		Req{Time: now.AddDate(0, 0, -1).UnixMilli(), Model: "old", Source: "s", Prompt: 999, Compl: 99, OK: true},
	)
	s := periodStats(t, path)

	day := s.PeriodSnapshot(PeriodDay, "", now)
	if day.Total.Reqs != 2 {
		t.Errorf("day window counted %d requests, want 2 (yesterday's record must be excluded)", day.Total.Reqs)
	}
	if day.Total.Prompt != 300 {
		t.Errorf("day prompt tokens = %d, want 300", day.Total.Prompt)
	}

	week := s.PeriodSnapshot(PeriodWeek, "", now)
	if week.Total.Reqs != 3 {
		t.Errorf("week window counted %d requests, want 3", week.Total.Reqs)
	}
	if week.Total.Prompt != 1299 {
		t.Errorf("week prompt tokens = %d, want 1299", week.Total.Prompt)
	}
}

// TestPeriodDimensionsAggregateIndependently checks the by_* rows actually
// split by their own key rather than all collapsing into one row.
func TestPeriodDimensionsAggregateIndependently(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
	auditPeriod(t, path,
		Req{Time: now.UnixMilli(), Key: "k1", Model: "alpha", Source: "s1", Prompt: 10, OK: true},
		Req{Time: now.UnixMilli(), Key: "k1", Model: "beta", Source: "s2", Prompt: 20, OK: true},
		Req{Time: now.UnixMilli(), Key: "k2", Model: "alpha", Source: "s1", Prompt: 30, OK: true},
	)
	s := periodStats(t, path)
	got := s.PeriodSnapshot(PeriodDay, "", now)

	if len(got.Models) != 2 {
		t.Fatalf("by_model has %d rows, want 2: %+v", len(got.Models), got.Models)
	}
	// alpha = 10 + 30
	for _, r := range got.Models {
		if r.Name == "alpha" && r.Prompt != 40 {
			t.Errorf("alpha prompt = %d, want 40", r.Prompt)
		}
	}
	if len(got.Srcs) != 2 {
		t.Errorf("by_source has %d rows, want 2: %+v", len(got.Srcs), got.Srcs)
	}
	if len(got.ByKey) != 2 {
		t.Errorf("by_key has %d rows, want 2: %+v", len(got.ByKey), got.ByKey)
	}
	// The rows must add up to the total, or the dashboard shows a total that
	// disagrees with its own table.
	var sum int64
	for _, r := range got.Models {
		sum += r.Prompt
	}
	if sum != got.Total.Prompt {
		t.Errorf("by_model prompts sum to %d but total is %d — the table would contradict the KPI", sum, got.Total.Prompt)
	}
}

// TestPeriodDayBucketsAreHourlyAndOrdered pins the timeline grain and its
// order. Day = hourly (a spike must be visible inside "today"); week/month =
// daily. Sorted chronologically, because a chart fed unsorted buckets draws
// nonsense.
func TestPeriodDayBucketsAreHourlyAndOrdered(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	now := time.Date(2026, 3, 11, 15, 0, 0, 0, time.UTC)
	// Deliberately out of order in the file.
	auditPeriod(t, path,
		Req{Time: now.Add(-2 * time.Hour).UnixMilli(), Model: "m", Prompt: 1, OK: true},
		Req{Time: now.Add(-4 * time.Hour).UnixMilli(), Model: "m", Prompt: 1, OK: true},
		Req{Time: now.UnixMilli(), Model: "m", Prompt: 1, OK: true},
	)
	s := periodStats(t, path)
	got := s.PeriodSnapshot(PeriodDay, "", now)
	if len(got.Bucket) != 3 {
		t.Fatalf("day buckets = %d, want 3 (hourly): %+v", len(got.Bucket), got.Bucket)
	}
	for i := 1; i < len(got.Bucket); i++ {
		if got.Bucket[i-1].Bucket >= got.Bucket[i].Bucket {
			t.Errorf("day buckets not chronological at %d: %q >= %q",
				i, got.Bucket[i-1].Bucket, got.Bucket[i].Bucket)
		}
	}
	// Weekly grain is daily, so the same three records collapse to one day.
	week := s.PeriodSnapshot(PeriodWeek, "", now)
	if len(week.Bucket) != 1 {
		t.Errorf("week buckets = %d, want 1 (daily grain): %+v", len(week.Bucket), week.Bucket)
	}
}

// TestPeriodKeyFilterIsolatesKeys guards the multi-tenant boundary: asking for
// one key must never return another key's rows in any dimension.
func TestPeriodKeyFilterIsolatesKeys(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
	auditPeriod(t, path,
		Req{Time: now.UnixMilli(), Key: "alice", Model: "m", Source: "s", Prompt: 10, OK: true},
		Req{Time: now.UnixMilli(), Key: "bob", Model: "m", Source: "s", Prompt: 999, OK: true},
	)
	s := periodStats(t, path)
	got := s.PeriodSnapshot(PeriodDay, "alice", now)
	if got.Total.Reqs != 1 || got.Total.Prompt != 10 {
		t.Errorf("alice's window = %d reqs / %d prompt, want 1 / 10 — another key leaked in",
			got.Total.Reqs, got.Total.Prompt)
	}
	if len(got.Srcs) != 1 || got.Srcs[0].Prompt != 10 {
		t.Errorf("by_source leaked: %+v", got.Srcs)
	}
}

// TestPeriodAllUsesLifetimeAggregates checks "all" still answers from the
// in-memory aggregates (no audit file needed) and carries the same rows the
// dashboard already renders.
func TestPeriodAllUsesLifetimeAggregates(t *testing.T) {
	s := NewStats(100)
	now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
	s.Record(Req{Time: now.AddDate(0, 0, -40).UnixMilli(), Key: "k", Model: "m", Source: "s", Prompt: 7, OK: true})
	s.Record(Req{Time: now.UnixMilli(), Key: "k", Model: "m", Source: "s", Prompt: 3, OK: true})

	got := s.PeriodSnapshot(PeriodAll, "", now)
	if got.Total.Reqs != 2 || got.Total.Prompt != 10 {
		t.Errorf("all view = %d reqs / %d prompt, want 2 / 10 (lifetime)", got.Total.Reqs, got.Total.Prompt)
	}
	if len(got.Models) != 1 || got.Models[0].Prompt != 10 {
		t.Errorf("all view by_model = %+v, want one row with 10 prompt", got.Models)
	}
	// A lifetime total has no timeline to draw; buckets must stay empty
	// rather than inventing 500 hourly points.
	if len(got.Bucket) != 0 {
		t.Errorf("all view produced %d buckets, want 0", len(got.Bucket))
	}
}

// TestPeriodEmptyWindowFlagsTruncated covers the "rotated away" case: a bounded
// window with nothing to show may mean "quiet day" or "the files are gone".
// The aggregator must not report a confidently wrong zero.
func TestPeriodEmptyWindowFlagsTruncated(t *testing.T) {
	dir := t.TempDir()
	s := periodStats(t, filepath.Join(dir, "audit.jsonl")) // file does not exist
	now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
	got := s.PeriodSnapshot(PeriodDay, "", now)
	if got.Total.Reqs != 0 {
		t.Errorf("empty window reqs = %d, want 0", got.Total.Reqs)
	}
	if !got.Truncated {
		t.Error("★ an empty bounded window must be flagged truncated — otherwise a " +
			"rotated-away week reads as 'no traffic that week'")
	}
}

// TestValidPeriodRejectsUnknown guards the deliberate strictness: a typo must
// be an error, not a silent fallback to lifetime totals.
func TestValidPeriodRejectsUnknown(t *testing.T) {
	for _, ok := range []Period{PeriodDay, PeriodWeek, PeriodMonth, PeriodAll} {
		if !ValidPeriod(ok) {
			t.Errorf("ValidPeriod(%q) = false, want true", ok)
		}
	}
	for _, bad := range []Period{"", "year", "hour", "today", "DAY "} {
		if ValidPeriod(bad) {
			t.Errorf("ValidPeriod(%q) = true, want false", bad)
		}
	}
}

// TestPeriodStatusDimensionCountsFailures covers by_status, which nothing else
// exercised — an aggregation branch with no test is exactly how a period view
// silently loses the failure pie. M8 (disabling the status dimension) is
// caught only because this test exists.
func TestPeriodStatusDimensionCountsFailures(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "audit.jsonl")
	now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
	auditPeriod(t, path,
		Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 200, OK: true, Prompt: 5},
		Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 502, OK: false, Prompt: 1},
		Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 502, OK: false, Prompt: 1},
		Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 429, OK: false, Prompt: 1},
	)
	s := periodStats(t, path)
	got := s.PeriodSnapshot(PeriodDay, "", now)

	byCode := map[string]int64{}
	for _, r := range got.Status {
		byCode[r.Name] = r.Reqs
	}
	if len(got.Status) != 3 {
		t.Fatalf("by_status has %d rows, want 3 (200/502/429): %+v", len(got.Status), got.Status)
	}
	if byCode["200"] != 1 || byCode["502"] != 2 || byCode["429"] != 1 {
		t.Errorf("by_status counts = %v, want 200:1 502:2 429:1", byCode)
	}
	// Status codes must sort numerically: "429" before "502" lexically would
	// render the failure legend in the wrong order.
	if got.Status[0].Name != "200" || got.Status[1].Name != "429" || got.Status[2].Name != "502" {
		t.Errorf("by_status not numerically sorted: %s, %s, %s",
			got.Status[0].Name, got.Status[1].Name, got.Status[2].Name)
	}
	// The failure rows must agree with the OK/Err split on the total, or the
	// pie and the KPI row contradict each other.
	if got.Total.OK != 1 || got.Total.Err != 3 {
		t.Errorf("total ok/err = %d/%d, want 1/3", got.Total.OK, got.Total.Err)
	}
}

func TestPeriodFromQuery(t *testing.T) {
	cases := []struct {
		query string
		want  Period
	}{
		{"", PeriodAll}, // no selector → unchanged legacy behaviour
		{"?period=day", PeriodDay},
		{"?period=WEEK", PeriodWeek},
		{"?period=month", PeriodMonth},
		{"?period=all", PeriodAll},
	}
	for _, c := range cases {
		r := httptest.NewRequest(http.MethodGet, "/api/stats"+c.query, nil)
		if got := periodFromQuery(r.URL.Query()); got != c.want {
			t.Errorf("periodFromQuery(%q) = %q, want %q", c.query, got, c.want)
		}
	}
}