package mcp import ( "context" "sync" "testing" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/yan1h/agent-coding-workflow/internal/audit" ) type captureRec struct { mu sync.Mutex entries []audit.Entry } func (r *captureRec) Record(_ context.Context, e audit.Entry) error { r.mu.Lock() defer r.mu.Unlock() r.entries = append(r.entries, e) return nil } func TestAuditWrap_InjectsTokenID(t *testing.T) { t.Parallel() cap := &captureRec{} wrapped := AuditWrap(cap) uid := uuid.New() ctx := withMcpTokenID(context.Background(), "tok-123") require.NoError(t, wrapped.Record(ctx, audit.Entry{ Action: "issue.create", UserID: &uid, Metadata: map[string]any{"slug": "demo"}, })) require.Len(t, cap.entries, 1) got := cap.entries[0] assert.Equal(t, true, got.Metadata["via_mcp"]) assert.Equal(t, "tok-123", got.Metadata["mcp_token_id"]) assert.Equal(t, "demo", got.Metadata["slug"]) } func TestAuditWrap_NoTokenID_Passthrough(t *testing.T) { t.Parallel() cap := &captureRec{} wrapped := AuditWrap(cap) require.NoError(t, wrapped.Record(context.Background(), audit.Entry{ Action: "user.login", })) require.Len(t, cap.entries, 1) got := cap.entries[0] _, hasViaMCP := got.Metadata["via_mcp"] assert.False(t, hasViaMCP) } func TestAuditWrap_NilInner_ReturnsNil(t *testing.T) { t.Parallel() assert.Nil(t, AuditWrap(nil)) } func TestAuditWrap_NilMetadata_Initialized(t *testing.T) { t.Parallel() cap := &captureRec{} wrapped := AuditWrap(cap) ctx := withMcpTokenID(context.Background(), "tok-x") require.NoError(t, wrapped.Record(ctx, audit.Entry{Action: "x"})) require.Len(t, cap.entries, 1) require.NotNil(t, cap.entries[0].Metadata) assert.Equal(t, "tok-x", cap.entries[0].Metadata["mcp_token_id"]) }