llmgw/internal/store/guard_test.go

package store

import (
	"testing"
	"time"

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

func TestAGuardPolicyIsAttachedOnceAndDetached(t *testing.T) {
	db := capacityStore(t)
	ctx := t.Context()
	a := guard.Attachment{Scope: "project", ScopeID: "p1", Policy: "finance"}
	for range 2 {
		if err := db.AttachGuard(ctx, a, "ilker"); err != nil {
			t.Fatal(err)
		}
	}
	if got, err := db.GuardAttachments(ctx); err != nil || len(got) != 1 || got[0] != a {
		t.Fatalf("%v %v", got, err)
	}
	if ok, err := db.DetachGuard(ctx, a); !ok || err != nil {
		t.Errorf("detach: %v %v", ok, err)
	}
	if ok, _ := db.DetachGuard(ctx, a); ok {
		t.Error("detached twice")
	}
}

func TestWhatTheGuardrailsDidCanBeReadBack(t *testing.T) {
	db := capacityStore(t)
	ctx := t.Context()
	now := time.Now()
	rows := []Usage{
		{At: now, KeyID: "k", ProjectID: "p1", Model: "m", ServedModel: "m", Status: 400, RequestID: "r1",
			GuardAction: "blocked", GuardPolicy: "base", GuardCategory: "secret", GuardMs: 3},
		{At: now, KeyID: "k", ProjectID: "p1", Model: "m", ServedModel: "m", Status: 200, RequestID: "r2",
			GuardAction: "blocked", GuardPolicy: "base", GuardCategory: "secret"},
		{At: now, KeyID: "k", ProjectID: "p2", Model: "m", ServedModel: "m", Status: 200, RequestID: "r3"},
	}
	if err := db.WriteUsage(ctx, rows); err != nil {
		t.Fatal(err)
	}
	ev, err := db.GuardEvents(ctx, now.Add(-time.Hour), 10)
	if err != nil || len(ev) != 2 || ev[0].Category != "secret" {
		t.Fatalf("%+v %v", ev, err)
	}
	c, err := db.GuardCounts(ctx, now.Add(-time.Hour))
	if err != nil || len(c) != 1 || c[0].Count != 2 || c[0].ProjectID != "p1" {
		t.Errorf("%+v %v", c, err)
	}
}