ModelRouter/internal/gateway/stats_unknown_key_test.go

package gateway

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

// TestSnapshotUnknownKeyDoesNotPanic covers a nil-dereference that killed the
// response for any /api/stats request carrying a ?key= the statistics had no
// entry for.
//
// The maps in Snapshot are keyed by the MASKED key id (keyID(req.Key)) — that
// is what by_key reports and what the dashboard sends. An admin may pass anything
// in the query parameter, including the full key, and `s.byKey[key]` then yields
// the zero value: a nil *Stat. rows() dereferenced every value, so the handler
// panicked and the client got an empty reply plus a panic stack in the log.
//
// Reproduced against the live gateway before the fix:
//
//	/api/stats?period=day&limit=1&key=sk-gw-7ffe12c5cbacec40bd39a7291df9a9ed
//	→ curl: (52) Empty reply from server
//	→ log: http: panic serving …: invalid memory address or nil pointer dereference
//	    llmsproxy/internal/gateway.rows  stats.go:771
//
// A key with no recorded traffic is a legitimate question, and zero is its
// answer — not a crash.
func TestSnapshotUnknownKeyDoesNotPanic(t *testing.T) {
	s := NewStats(16)
	// One recorded key, masked the way Record masks it.
	s.Record(Req{Key: "***aaaaaa", Type: "chat", Model: "m", Prompt: 5, Compl: 2})

	cases := []struct {
		name string
		key  string
	}{
		{"masked key that exists", "***aaaaaa"},
		{"masked key that does not exist", "***zzzzzz"},
		{"full key rather than the mask", "sk-gw-full-key-value"},
		{"arbitrary text", "not-a-key-at-all"},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			// A panic here fails the subtest by unwinding through the test
			// goroutine; rows() used to take the whole process with it.
			out := s.Snapshot(10, c.key)
			byKey, ok := out["by_key"].([]StatsRow)
			if !ok {
				t.Fatalf("by_key is %T, want []StatsRow", out["by_key"])
			}
			total, _ := out["total"].(Stat)
			t.Logf("key=%q → by_key %d 条, total.reqs=%d", c.key, len(byKey), total.Reqs)
		})
	}
}

// TestSnapshotUnknownKeyReportsZeroNotAbsence pins the semantics: a key with no
// traffic answers zero, and the by_key row still exists so a dashboard that keys
// off its presence renders a row of zeros rather than dropping the key.
func TestSnapshotUnknownKeyReportsZeroNotAbsence(t *testing.T) {
	s := NewStats(16)
	out := s.Snapshot(10, "***zzzzzz")
	byKey, _ := out["by_key"].([]StatsRow)
	if len(byKey) != 1 {
		t.Fatalf("by_key has %d rows, want 1: the queried key must appear with zeros", len(byKey))
	}
	if byKey[0].Name != "***zzzzzz" {
		t.Errorf("row name = %q, want ***zzzzzz", byKey[0].Name)
	}
	if byKey[0].Reqs != 0 {
		t.Errorf("reqs = %d, want 0", byKey[0].Reqs)
	}
}

// TestPeriodSnapshotUnknownKeyDoesNotPanic is the same guard on the period view,
// which builds its own buckets and calls the same flattening helpers.
func TestPeriodSnapshotUnknownKeyDoesNotPanic(t *testing.T) {
	dir := t.TempDir()
	audit := dir + "/audit.jsonl"
	writeAuditFile(t, audit, 3, "***aaaaaa", time.Now().Add(-time.Hour).UnixMilli())
	s := periodStats(t, audit)

	for _, key := range []string{"***aaaaaa", "***zzzzzz", "sk-gw-full-key-value"} {
		out := s.PeriodSnapshot(PeriodAll, key, time.Now())
		t.Logf("key=%q → total.reqs=%d byKey=%d buckets=%d",
			key, out.Total.Reqs, len(out.ByKey), len(out.Bucket))
	}
}

// TestRowsSkipsNilValues pins the defensive behaviour directly, so a future
// caller that reintroduces a nil entry fails here rather than in production.
func TestRowsSkipsNilValues(t *testing.T) {
	in := map[string]*Stat{
		"has-traffic": {Reqs: 5},
		"no-entry":    nil,
	}
	got := rows(in)
	if len(got) != 1 {
		t.Fatalf("rows() returned %d rows, want 1 (the nil entry must be skipped): %+v", len(got), got)
	}
	if got[0].Name != "has-traffic" {
		t.Errorf("kept %q, want has-traffic", got[0].Name)
	}
}

// TestRowsEmptyIsEmptyNotNil keeps the JSON shape honest: an all-nil map yields
// an empty list, which encodes as [] rather than null — the frontend does
// Array.isArray checks on these fields.
func TestRowsEmptyIsEmptyNotNil(t *testing.T) {
	got := rows(map[string]*Stat{"a": nil})
	if got == nil {
		t.Fatal("rows() returned nil; the report field must stay [] instead of null")
	}
	if len(got) != 0 {
		t.Fatalf("got %d rows, want 0", len(got))
	}
}

// TestStatsAPIUnknownKeyReturnsJSON exercises the HTTP path so the assertion
// covers the reported symptom (an empty reply), not just the helper.
func TestStatsAPIUnknownKeyReturnsJSON(t *testing.T) {
	g := newTestGateway(t)
	srv := httptest.NewServer(g.Handler())
	defer srv.Close()

	req, err := http.NewRequest("GET", srv.URL+"/api/stats?period=day&limit=1&key=sk-gw-full-key", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer sk-test")
	resp, err := srv.Client().Do(req)
	if err != nil {
		t.Fatalf("request failed (an empty reply is what the live gateway produced): %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200", resp.StatusCode)
	}
}