llmgw/internal/control/adminapi/guard_test.go

package adminapi

import (
	"context"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

type memGuard struct{ att []guard.Attachment }

func (m *memGuard) GuardAttachments(context.Context) ([]guard.Attachment, error) { return m.att, nil }
func (m *memGuard) AttachGuard(_ context.Context, a guard.Attachment, _ string) error {
	if !slices.Contains(m.att, a) {
		m.att = append(m.att, a)
	}
	return nil
}
func (m *memGuard) DetachGuard(_ context.Context, a guard.Attachment) (bool, error) {
	n := len(m.att)
	m.att = slices.DeleteFunc(m.att, func(x guard.Attachment) bool { return x == a })
	return len(m.att) < n, nil
}

func TestGuardPoliciesAreAttachedByTarget(t *testing.T) {
	a, _ := setup(t, testModels())
	if code := call(t, a, "GET", "/api/v1/guard", "", nil); code != 501 {
		t.Errorf("without policies: %d", code)
	}
	st := &memGuard{}
	invalidated := 0
	a.SetGuard(&GuardAdmin{
		Library: func() guard.Config {
			return guard.Config{Default: []string{"base"}, Policies: map[string]guard.Policy{"base": {}, "finance": {}}}
		},
		Store: st, Invalidate: func() { invalidated++ },
	})
	if code := call(t, a, "POST", "/api/v1/guard/attachments", `{"target":"key:k1","policy":"finance"}`, nil); code != 200 {
		t.Fatalf("attach: %d", code)
	}
	if invalidated != 1 || len(st.att) != 1 || st.att[0] != (guard.Attachment{Scope: "key", ScopeID: "k1", Policy: "finance"}) {
		t.Errorf("stored %v, invalidated %d", st.att, invalidated)
	}
	for _, bad := range []string{`{"target":"key:k1","policy":"nope"}`, `{"target":"galaxy:x","policy":"finance"}`, `{"target":"key:","policy":"finance"}`} {
		if code := call(t, a, "POST", "/api/v1/guard/attachments", bad, nil); code != 400 {
			t.Errorf("%s: %d", bad, code)
		}
	}
	var v GuardView
	if code := call(t, a, "GET", "/api/v1/guard", "", &v); code != 200 || len(v.Attachments) != 1 ||
		!slices.Equal(v.Policies, []string{"base", "finance"}) || !slices.Equal(v.Default, []string{"base"}) {
		t.Errorf("%d %+v", code, v)
	}
	if code := call(t, a, "DELETE", "/api/v1/guard/attachments", `{"target":"key:k1","policy":"finance"}`, nil); code != 200 || len(st.att) != 0 {
		t.Errorf("detach: %d %v", code, st.att)
	}
	if code := call(t, a, "DELETE", "/api/v1/guard/attachments", `{"target":"key:k1","policy":"finance"}`, nil); code != 404 {
		t.Errorf("detach again: %d", code)
	}
}

type memAudit struct{ f store.AuditFilter }

func (m *memAudit) AuditEvents(_ context.Context, f store.AuditFilter) ([]store.AuditEvent, error) {
	m.f = f
	return []store.AuditEvent{{Actor: "ilker", Action: "POST /api/v1/models"}}, nil
}

func TestTheAuditLogIsReadWithFilters(t *testing.T) {
	a, _ := setup(t, testModels())
	m := &memAudit{}
	a.SetAudit(m)
	var ev []store.AuditEvent
	if code := call(t, a, "GET", "/api/v1/audit?since=2d&actor=ilker&contains=models&limit=5", "", &ev); code != 200 || len(ev) != 1 {
		t.Fatalf("%d %v", code, ev)
	}
	if m.f.Actor != "ilker" || m.f.Contains != "models" || m.f.Limit != 5 || time.Since(m.f.Since) < 47*time.Hour {
		t.Errorf("filter %+v", m.f)
	}
	if code := call(t, a, "GET", "/api/v1/audit?since=soon", "", nil); code != 400 {
		t.Errorf("a bad since: %d", code)
	}
}

type memUsageExport struct {
	memAudit
	by     []string
	period string
}

func (m *memUsageExport) UsageByScope(context.Context, store.Scope, time.Time) ([]store.ScopeUsage, error) {
	return nil, nil
}

func (m *memUsageExport) UsageExport(_ context.Context, from, to time.Time, period string, by []string) ([]store.ExportRow, error) {
	m.by, m.period = by, period
	return []store.ExportRow{{Period: from, Project: "=HYPERLINK(evil)", Model: "m", Requests: 3, CostMicros: 1500000}}, nil
}

func TestUsageIsExportedAsCSVThatASpreadsheetCannotRun(t *testing.T) {
	a, _ := setup(t, testModels())
	m := &memUsageExport{}
	a.usage = m
	mux := http.NewServeMux()
	a.Register(mux)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, httptest.NewRequest("GET", "/api/v1/usage/export?from=2026-09-01&to=2026-10-01&period=month&by=project,model", nil))
	body := rec.Body.String()
	if rec.Code != 200 || !strings.HasPrefix(rec.Header().Get("Content-Type"), "text/csv") ||
		!strings.Contains(body, "period_start,org,team,project") || !strings.Contains(body, "'=HYPERLINK") ||
		!strings.Contains(body, "1.500000") || m.period != "month" || !slices.Equal(m.by, []string{"project", "model"}) {
		t.Errorf("%d %s %+v", rec.Code, body, m)
	}
}

// memGuardHistory is a guard store that also remembers what was done.
type memGuardHistory struct {
	memGuard
	since time.Time
}

func (m *memGuardHistory) GuardEvents(_ context.Context, since time.Time, _ int) ([]store.GuardEvent, error) {
	m.since = since
	return []store.GuardEvent{{RequestID: "r1", Model: "m", Status: 400, Action: "block", Policy: "base", Category: "S1", GuardMs: 410}}, nil
}

func (m *memGuardHistory) GuardCounts(context.Context, time.Time) ([]store.GuardCount, error) {
	return []store.GuardCount{
		{Action: "flag", Policy: "base", Category: "email", Count: 2},
		{Action: "block", Policy: "base", Category: "S1", Count: 7},
	}, nil
}

func TestTheGuardrailsSayWhatEachPolicyDoesAndWhatTheyDid(t *testing.T) {
	a, _ := setup(t, testModels())
	h := &memGuardHistory{}
	lib := guard.Config{Policies: map[string]guard.Policy{
		"base": {Fail: "closed", Input: guard.Rules{PII: map[string]string{"email": "redact"}, Injection: "block"},
			Judge: &guard.JudgePolicy{Model: "llama-guard-3-1b", Input: "block", Categories: []string{"S1", "S9"}}},
	}}
	a.SetGuard(&GuardAdmin{Store: h, Library: func() guard.Config { return lib }})

	var v GuardView
	if code := call(t, a, "GET", "/api/v1/guard", "", &v); code != 200 || len(v.Details) != 1 {
		t.Fatalf("%d %+v", code, v)
	}
	d := v.Details[0]
	want := []string{"input: pii email → redact", "input: prompt injection → block", "judge llama-guard-3-1b: input → block (S1 S9)"}
	if d.Mode != "enforce" || d.Fail != "closed" || d.Judge != "llama-guard-3-1b" || !slices.Equal(d.Checks, want) {
		t.Errorf("%+v", d)
	}

	var ev GuardEventsView
	if code := call(t, a, "GET", "/api/v1/guard/events?since=1h", "", &ev); code != 200 {
		t.Fatalf("events: %d", code)
	}
	if time.Since(h.since) > time.Hour+time.Minute || time.Since(h.since) < time.Hour-time.Minute {
		t.Errorf("asked since %s", h.since)
	}
	if len(ev.Counts) != 2 || ev.Counts[0].Count != 7 || len(ev.Events) != 1 || ev.Events[0].GuardMs != 410 {
		t.Errorf("most first, and the events: %+v", ev)
	}
	if code := call(t, a, "GET", "/api/v1/guard/events?since=soon", "", nil); code != 400 {
		t.Errorf("a bad window: %d", code)
	}
}