Hein
b17241b928
Some checks failed
CI / build-and-test (push) Failing after -31m49s
* Add tool tracking to AccessTracker and metrics * Update tests to validate tool tracking functionality * Modify middleware to record tool usage * Enhance observability with tool context * Update UI to display unique tools in metrics
97 lines
3.5 KiB
Go
97 lines
3.5 KiB
Go
package auth
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAccessTrackerRecordAndSnapshot(t *testing.T) {
|
|
tracker := NewAccessTracker()
|
|
older := time.Date(2026, 4, 4, 10, 0, 0, 0, time.UTC)
|
|
newer := older.Add(2 * time.Minute)
|
|
|
|
tracker.Record("client-a", "/files", "10.0.0.1:1234", "agent-a", "", older)
|
|
tracker.Record("client-b", "/mcp", "10.0.0.2:1234", "agent-b", "list_projects", newer)
|
|
tracker.Record("client-a", "/files/1", "10.0.0.1:1234", "agent-a2", "", newer.Add(30*time.Second))
|
|
|
|
snap := tracker.Snapshot()
|
|
if len(snap) != 2 {
|
|
t.Fatalf("len(snapshot) = %d, want 2", len(snap))
|
|
}
|
|
if snap[0].KeyID != "client-a" {
|
|
t.Fatalf("snapshot[0].KeyID = %q, want client-a", snap[0].KeyID)
|
|
}
|
|
if snap[0].RequestCount != 2 {
|
|
t.Fatalf("snapshot[0].RequestCount = %d, want 2", snap[0].RequestCount)
|
|
}
|
|
if snap[0].LastPath != "/files/1" {
|
|
t.Fatalf("snapshot[0].LastPath = %q, want /files/1", snap[0].LastPath)
|
|
}
|
|
if snap[0].UserAgent != "agent-a2" {
|
|
t.Fatalf("snapshot[0].UserAgent = %q, want agent-a2", snap[0].UserAgent)
|
|
}
|
|
if snap[0].RemoteAddr != "10.0.0.1" {
|
|
t.Fatalf("snapshot[0].RemoteAddr = %q, want 10.0.0.1", snap[0].RemoteAddr)
|
|
}
|
|
}
|
|
|
|
func TestAccessTrackerConnectedCount(t *testing.T) {
|
|
tracker := NewAccessTracker()
|
|
now := time.Date(2026, 4, 4, 12, 0, 0, 0, time.UTC)
|
|
|
|
tracker.Record("recent", "/mcp", "", "", "", now.Add(-2*time.Minute))
|
|
tracker.Record("stale", "/mcp", "", "", "", now.Add(-11*time.Minute))
|
|
|
|
if got := tracker.ConnectedCount(now, 10*time.Minute); got != 1 {
|
|
t.Fatalf("ConnectedCount() = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestAccessTrackerMetrics(t *testing.T) {
|
|
tracker := NewAccessTracker()
|
|
now := time.Date(2026, 4, 4, 12, 0, 0, 0, time.UTC)
|
|
|
|
tracker.Record("client-a", "/mcp", "10.0.0.1:1234", "agent-a", "list_projects", now)
|
|
tracker.Record("client-a", "/mcp", "10.0.0.1:1234", "agent-a", "list_projects", now.Add(1*time.Second))
|
|
tracker.Record("client-b", "/files", "10.0.0.2:5678", "agent-b", "", now.Add(2*time.Second))
|
|
tracker.Record("client-c", "/files", "10.0.0.2:5678", "agent-b", "search_thoughts", now.Add(3*time.Second))
|
|
|
|
metrics := tracker.Metrics(5)
|
|
if metrics.TotalRequests != 4 {
|
|
t.Fatalf("TotalRequests = %d, want 4", metrics.TotalRequests)
|
|
}
|
|
if metrics.UniquePrincipals != 3 {
|
|
t.Fatalf("UniquePrincipals = %d, want 3", metrics.UniquePrincipals)
|
|
}
|
|
if metrics.UniqueIPs != 2 {
|
|
t.Fatalf("UniqueIPs = %d, want 2", metrics.UniqueIPs)
|
|
}
|
|
if metrics.UniqueAgents != 2 {
|
|
t.Fatalf("UniqueAgents = %d, want 2", metrics.UniqueAgents)
|
|
}
|
|
if metrics.UniqueTools != 2 {
|
|
t.Fatalf("UniqueTools = %d, want 2", metrics.UniqueTools)
|
|
}
|
|
if len(metrics.TopIPs) != 2 {
|
|
t.Fatalf("len(TopIPs) = %d, want 2", len(metrics.TopIPs))
|
|
}
|
|
if metrics.TopIPs[0].RequestCount != 2 || metrics.TopIPs[1].RequestCount != 2 {
|
|
t.Fatalf("TopIPs counts = %+v, want both counts to be 2", metrics.TopIPs)
|
|
}
|
|
if metrics.TopIPs[0].Key != "10.0.0.1" && metrics.TopIPs[0].Key != "10.0.0.2" {
|
|
t.Fatalf("TopIPs[0].Key = %q, want normalized IP", metrics.TopIPs[0].Key)
|
|
}
|
|
if len(metrics.TopAgents) != 2 {
|
|
t.Fatalf("len(TopAgents) = %d, want 2", len(metrics.TopAgents))
|
|
}
|
|
if metrics.TopAgents[0].RequestCount != 2 || metrics.TopAgents[1].RequestCount != 2 {
|
|
t.Fatalf("TopAgents counts = %+v, want both counts to be 2", metrics.TopAgents)
|
|
}
|
|
if len(metrics.TopTools) != 2 {
|
|
t.Fatalf("len(TopTools) = %d, want 2", len(metrics.TopTools))
|
|
}
|
|
if metrics.TopTools[0].Key != "list_projects" || metrics.TopTools[0].RequestCount != 2 {
|
|
t.Fatalf("TopTools[0] = %+v, want list_projects with count 2", metrics.TopTools[0])
|
|
}
|
|
}
|