From 1f74b227b960eabc2922e8908aebea0b9e16a513 Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Thu, 19 Feb 2026 14:48:46 +0000 Subject: [PATCH] test: add comprehensive tests for ITR commands, DAG, RLM, and SecureBus - ITR command serialization round-trip tests - DAG planner, resolver, router, and replan tests - RLM fanout concurrency and strategy planning tests - SecureBus audit log and policy enforcement tests --- pkg/itr/commands_test.go | 228 ++++++++++++++++++++++++++ pkg/itr/dag/planner_test.go | 184 +++++++++++++++++++++ pkg/itr/dag/replan_test.go | 18 ++ pkg/itr/dag/resolver_test.go | 120 ++++++++++++++ pkg/itr/dag/router_test.go | 60 +++++++ pkg/rlm/fanout_test.go | 130 +++++++++++++++ pkg/rlm/strategy_test.go | 86 ++++++++++ pkg/security/securebus/audit_test.go | 132 +++++++++++++++ pkg/security/securebus/policy_test.go | 103 ++++++++++++ 9 files changed, 1061 insertions(+) create mode 100644 pkg/itr/commands_test.go create mode 100644 pkg/itr/dag/planner_test.go create mode 100644 pkg/itr/dag/replan_test.go create mode 100644 pkg/itr/dag/resolver_test.go create mode 100644 pkg/itr/dag/router_test.go create mode 100644 pkg/rlm/fanout_test.go create mode 100644 pkg/rlm/strategy_test.go create mode 100644 pkg/security/securebus/audit_test.go create mode 100644 pkg/security/securebus/policy_test.go diff --git a/pkg/itr/commands_test.go b/pkg/itr/commands_test.go new file mode 100644 index 000000000..ed9712ae7 --- /dev/null +++ b/pkg/itr/commands_test.go @@ -0,0 +1,228 @@ +package itr + +import ( + jsonv2 "github.com/go-json-experiment/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestToolRequestMarshalRoundTrip(t *testing.T) { + tests := []struct { + name string + req ToolRequest + }{ + { + name: "ToolExec", + req: NewToolExecRequest("id-1", "sess-1", "tc-1", "read_file", `{"path":"/tmp/a.txt"}`), + }, + { + name: "Peek", + req: NewPeekRequest("id-2", "sess-1", 0, 100, 500), + }, + { + name: "Grep", + req: NewGrepRequest("id-3", "sess-1", 1, "func main", 10, true), + }, + { + name: "Partition", + req: NewPartitionRequest("id-4", "sess-1", 2, 8, "semantic", 128, true), + }, + { + name: "Recurse", + req: NewRecurseRequest("id-5", "sess-1", 1, "summarize this", "ctx-key-99", 3), + }, + { + name: "Final", + req: NewFinalRequest("id-6", "sess-1", 0, "The answer is 42.", "final_ans"), + }, + { + name: "ToolSearch", + req: NewToolSearchRequest("id-7", "sess-1", "file operations", 5), + }, + { + name: "CodeExec", + req: NewCodeExecRequest("id-8", "sess-1", "print('hello')", "python-wasm"), + }, + { + name: "DAGPlan", + req: NewDAGPlanRequest("id-9", "sess-1", DAGPlan{ + Nodes: []DAGNode{ + {ID: "n1", Type: CmdToolExec, Payload: ToolExec{ToolName: "search", ArgsJSON: `{"q":"test"}`}}, + {ID: "n2", Type: CmdToolExec, Payload: ToolExec{ToolName: "read", ArgsJSON: `{"path":"#noden1"}`}, DependsOn: []string{"n1"}}, + }, + MaxParallel: 4, + TokenBudget: 50000, + JoinerQuery: "Synthesize the results", + }), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := tt.req.Marshal() + require.NoError(t, err) + require.NotEmpty(t, data) + + decoded, err := UnmarshalRequest(data) + require.NoError(t, err) + + assert.Equal(t, tt.req.ID, decoded.ID) + assert.Equal(t, tt.req.Type, decoded.Type) + assert.Equal(t, tt.req.SessionKey, decoded.SessionKey) + assert.Equal(t, tt.req.Depth, decoded.Depth) + assert.Equal(t, tt.req.ToolCallID, decoded.ToolCallID) + }) + } +} + +func TestToolExecPayloadPreservation(t *testing.T) { + req := NewToolExecRequest("id-1", "s", "tc", "shell", `{"cmd":"ls -la"}`) + data, err := req.Marshal() + require.NoError(t, err) + + decoded, err := UnmarshalRequest(data) + require.NoError(t, err) + + payload, ok := decoded.Payload.(ToolExec) + require.True(t, ok, "payload should be ToolExec") + assert.Equal(t, "shell", payload.ToolName) + assert.Equal(t, `{"cmd":"ls -la"}`, payload.ArgsJSON) +} + +func TestGrepPayloadPreservation(t *testing.T) { + req := NewGrepRequest("g1", "s", 2, "error.*fatal", 25, true) + data, err := req.Marshal() + require.NoError(t, err) + + decoded, err := UnmarshalRequest(data) + require.NoError(t, err) + + payload, ok := decoded.Payload.(Grep) + require.True(t, ok) + assert.Equal(t, "error.*fatal", payload.Pattern) + assert.Equal(t, uint32(25), payload.MaxMatches) + assert.True(t, payload.CaseInsensitive) +} + +func TestDAGPlanPayloadPreservation(t *testing.T) { + plan := DAGPlan{ + Nodes: []DAGNode{ + {ID: "a", Type: CmdToolSearch, Payload: ToolSearch{Query: "files", MaxResults: 5}, DependsOn: nil}, + {ID: "b", Type: CmdToolExec, Payload: ToolExec{ToolName: "read", ArgsJSON: `{"path":"#nodea"}`}, DependsOn: []string{"a"}}, + }, + MaxParallel: 2, + JoinerQuery: "combine results", + } + req := NewDAGPlanRequest("d1", "s", plan) + data, err := req.Marshal() + require.NoError(t, err) + + decoded, err := UnmarshalRequest(data) + require.NoError(t, err) + + dagPlan, ok := decoded.Payload.(DAGPlan) + require.True(t, ok) + assert.Len(t, dagPlan.Nodes, 2) + assert.Equal(t, "a", dagPlan.Nodes[0].ID) + assert.Equal(t, CmdToolSearch, dagPlan.Nodes[0].Type) + assert.Equal(t, []string{"a"}, dagPlan.Nodes[1].DependsOn) + assert.Equal(t, uint8(2), dagPlan.MaxParallel) + assert.Equal(t, "combine results", dagPlan.JoinerQuery) + + ts, ok := dagPlan.Nodes[0].Payload.(ToolSearch) + require.True(t, ok, "ToolSearch payload should be retyped after unmarshal") + assert.Equal(t, "files", ts.Query) + assert.Equal(t, uint8(5), ts.MaxResults) + + te, ok := dagPlan.Nodes[1].Payload.(ToolExec) + require.True(t, ok, "ToolExec payload should be retyped after unmarshal") + assert.Equal(t, "read", te.ToolName) + assert.Contains(t, te.ArgsJSON, "#nodea") +} + +func TestToolResponseMarshalRoundTrip(t *testing.T) { + tests := []struct { + name string + resp ToolResponse + }{ + { + name: "success", + resp: NewSuccessResponse("r1", `{"content":"hello"}`, 150), + }, + { + name: "error", + resp: NewErrorResponse("r2", "tool not found: badtool"), + }, + { + name: "leak", + resp: NewLeakResponse("r3", "redacted content", []string{"api_key", "token"}), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := tt.resp.Marshal() + require.NoError(t, err) + + decoded, err := UnmarshalResponse(data) + require.NoError(t, err) + + assert.Equal(t, tt.resp.ID, decoded.ID) + assert.Equal(t, tt.resp.Result, decoded.Result) + assert.Equal(t, tt.resp.IsError, decoded.IsError) + assert.Equal(t, tt.resp.LeakDetected, decoded.LeakDetected) + assert.Equal(t, tt.resp.CostTokens, decoded.CostTokens) + assert.Equal(t, tt.resp.RedactedKeys, decoded.RedactedKeys) + }) + } +} + +func TestUnmarshalRequestJSON_UnknownType(t *testing.T) { + data, _ := jsonv2.Marshal(map[string]interface{}{ + "id": "bad", + "type": "nonexistent_command", + }) + _, err := UnmarshalRequestJSON(data) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown command type") +} + +func TestUnmarshalRequestJSON_InvalidJSON(t *testing.T) { + _, err := UnmarshalRequestJSON([]byte(`{invalid`)) + assert.Error(t, err) +} + +func TestMarshalRequestFB_UnknownType(t *testing.T) { + _, err := MarshalRequestFB(ToolRequest{ID: "bad", Type: CommandType("bogus")}) + assert.Error(t, err) +} + +func TestUnmarshalRequestFB_Garbage(t *testing.T) { + _, err := UnmarshalRequestFB([]byte{0, 0, 0, 0}) + assert.Error(t, err) +} + +func TestRequestJSON_Roundtrip(t *testing.T) { + orig := NewToolExecRequest("j1", "s", "tc", "shell", `{"cmd":"ls"}`) + data, err := jsonv2.Marshal(orig) + require.NoError(t, err) + + decoded, err := UnmarshalRequestJSON(data) + require.NoError(t, err) + assert.Equal(t, orig.ID, decoded.ID) + p := decoded.Payload.(ToolExec) + assert.Equal(t, "shell", p.ToolName) +} + +func TestResponseJSON_Roundtrip(t *testing.T) { + orig := NewSuccessResponse("j2", "ok", 10) + data, err := jsonv2.Marshal(orig) + require.NoError(t, err) + + decoded, err := UnmarshalResponseJSON(data) + require.NoError(t, err) + assert.Equal(t, orig.ID, decoded.ID) + assert.Equal(t, "ok", decoded.Result) +} diff --git a/pkg/itr/dag/planner_test.go b/pkg/itr/dag/planner_test.go new file mode 100644 index 000000000..0ff511978 --- /dev/null +++ b/pkg/itr/dag/planner_test.go @@ -0,0 +1,184 @@ +package dag + +import ( + "context" + jsonv2 "github.com/go-json-experiment/json" + "testing" + + "github.com/sipeed/picoclaw/pkg/itr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExtractJSON_PlainJSON(t *testing.T) { + input := `{"nodes": [{"id": "n1"}]}` + assert.Equal(t, input, extractJSON(input)) +} + +func TestExtractJSON_MarkdownFenced(t *testing.T) { + input := "Here is the plan:\n```json\n{\"nodes\": [{\"id\": \"n1\"}]}\n```\nDone." + assert.Equal(t, `{"nodes": [{"id": "n1"}]}`, extractJSON(input)) +} + +func TestExtractJSON_GenericFenced(t *testing.T) { + input := "```\n{\"nodes\": []}\n```" + assert.Equal(t, `{"nodes": []}`, extractJSON(input)) +} + +func TestExtractJSON_LeadingText(t *testing.T) { + input := "The plan is: {\"nodes\":[]}" + assert.Equal(t, `{"nodes":[]}`, extractJSON(input)) +} + +func TestFindIndex(t *testing.T) { + assert.Equal(t, 0, findIndex("abc", "a")) + assert.Equal(t, 2, findIndex("abc", "c")) + assert.Equal(t, -1, findIndex("abc", "z")) + assert.Equal(t, -1, findIndex("", "a")) + assert.Equal(t, -1, findIndex("ab", "abc")) +} + +func TestValidatePlan_Valid(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "a", Type: itr.CmdToolExec}, + {ID: "b", Type: itr.CmdToolExec, DependsOn: []string{"a"}}, + }, + } + assert.NoError(t, validatePlan(plan)) +} + +func TestValidatePlan_Empty(t *testing.T) { + plan := &itr.DAGPlan{Nodes: nil} + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "no nodes") +} + +func TestValidatePlan_DuplicateID(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "x", Type: itr.CmdToolExec}, + {ID: "x", Type: itr.CmdToolExec}, + }, + } + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "duplicate node ID") +} + +func TestValidatePlan_EmptyID(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{{ID: "", Type: itr.CmdToolExec}}, + } + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty ID") +} + +func TestValidatePlan_UnknownDependency(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "a", Type: itr.CmdToolExec, DependsOn: []string{"missing"}}, + }, + } + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown node") +} + +func TestValidatePlan_SelfDependency(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "a", Type: itr.CmdToolExec, DependsOn: []string{"a"}}, + }, + } + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "depends on itself") +} + +func TestValidatePlan_CyclicDependency(t *testing.T) { + plan := &itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "a", Type: itr.CmdToolExec, DependsOn: []string{"b"}}, + {ID: "b", Type: itr.CmdToolExec, DependsOn: []string{"a"}}, + }, + } + err := validatePlan(plan) + require.Error(t, err) + assert.Contains(t, err.Error(), "cycle detected") +} + +func TestParsePlanResponse_ValidJSON(t *testing.T) { + input := `{ + "nodes": [ + {"id": "n1", "type": "tool_exec", "payload": {"tool_name": "read_file", "args_json": "{\"path\":\"/tmp/a\"}"}} + ], + "joiner_query": "summarize" + }` + + plan, err := parsePlanResponse(input) + require.NoError(t, err) + require.Len(t, plan.Nodes, 1) + assert.Equal(t, "n1", plan.Nodes[0].ID) + assert.Equal(t, itr.CmdToolExec, plan.Nodes[0].Type) + assert.Equal(t, "summarize", plan.JoinerQuery) + + te, ok := plan.Nodes[0].Payload.(itr.ToolExec) + require.True(t, ok) + assert.Equal(t, "read_file", te.ToolName) +} + +func TestParsePlanResponse_WithMarkdownFence(t *testing.T) { + input := "```json\n" + `{"nodes": [{"id": "x", "type": "tool_search", "payload": {"query": "files"}}]}` + "\n```" + + plan, err := parsePlanResponse(input) + require.NoError(t, err) + require.Len(t, plan.Nodes, 1) + + ts, ok := plan.Nodes[0].Payload.(itr.ToolSearch) + require.True(t, ok) + assert.Equal(t, "files", ts.Query) +} + +func TestParsePlanResponse_InvalidJSON(t *testing.T) { + _, err := parsePlanResponse("not json at all") + assert.Error(t, err) +} + +func TestPlannerPlanE2E(t *testing.T) { + mockLLM := func(ctx context.Context, systemPrompt, userQuery string) (string, uint32, error) { + plan := itr.DAGPlan{ + Nodes: []itr.DAGNode{ + {ID: "search", Type: itr.CmdToolExec, + Payload: itr.ToolExec{ToolName: "search", ArgsJSON: `{"q":"test"}`}}, + {ID: "read", Type: itr.CmdToolExec, + Payload: itr.ToolExec{ToolName: "read_file", ArgsJSON: `{"path":"#nodesearch"}`}, + DependsOn: []string{"search"}}, + }, + JoinerQuery: "Combine results", + } + b, _ := jsonv2.Marshal(plan) + return string(b), 100, nil + } + + planner := NewPlanner(mockLLM, nil, DefaultPlannerConfig()) + plan, tokens, err := planner.Plan(context.Background(), "search and read", nil) + require.NoError(t, err) + assert.Equal(t, uint32(100), tokens) + require.Len(t, plan.Nodes, 2) + assert.Equal(t, "search", plan.Nodes[0].ID) + assert.Equal(t, uint8(8), plan.MaxParallel) +} + +func TestPlannerPlanLLMError(t *testing.T) { + mockLLM := func(ctx context.Context, systemPrompt, userQuery string) (string, uint32, error) { + return "", 50, assert.AnError + } + + planner := NewPlanner(mockLLM, nil, DefaultPlannerConfig()) + _, tokens, err := planner.Plan(context.Background(), "anything", nil) + assert.Error(t, err) + assert.Equal(t, uint32(50), tokens) +} diff --git a/pkg/itr/dag/replan_test.go b/pkg/itr/dag/replan_test.go new file mode 100644 index 000000000..c0358c145 --- /dev/null +++ b/pkg/itr/dag/replan_test.go @@ -0,0 +1,18 @@ +package dag + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNeedsReplan(t *testing.T) { + assert.True(t, needsReplan("The task is incomplete [NEEDS_MORE_STEPS]")) + assert.True(t, needsReplan("[NEEDS_MORE_STEPS]")) + assert.False(t, needsReplan("Task complete. Here is the answer.")) + assert.False(t, needsReplan("")) +} + +func TestReplanSentinelIsConsistent(t *testing.T) { + assert.Equal(t, "[NEEDS_MORE_STEPS]", replanSentinel) +} diff --git a/pkg/itr/dag/resolver_test.go b/pkg/itr/dag/resolver_test.go new file mode 100644 index 000000000..da3cc7697 --- /dev/null +++ b/pkg/itr/dag/resolver_test.go @@ -0,0 +1,120 @@ +package dag + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTopologicalOrderLinear(t *testing.T) { + states := map[string]*nodeState{ + "a": newNodeState("a", nil), + "b": newNodeState("b", []string{"a"}), + "c": newNodeState("c", []string{"b"}), + } + + waves, err := topologicalOrder(states) + require.NoError(t, err) + require.Len(t, waves, 3) + assert.Equal(t, []string{"a"}, waves[0]) + assert.Equal(t, []string{"b"}, waves[1]) + assert.Equal(t, []string{"c"}, waves[2]) +} + +func TestTopologicalOrderParallel(t *testing.T) { + states := map[string]*nodeState{ + "a": newNodeState("a", nil), + "b": newNodeState("b", nil), + "c": newNodeState("c", []string{"a", "b"}), + } + + waves, err := topologicalOrder(states) + require.NoError(t, err) + require.Len(t, waves, 2) + + assert.Len(t, waves[0], 2) + assert.Contains(t, waves[0], "a") + assert.Contains(t, waves[0], "b") + assert.Equal(t, []string{"c"}, waves[1]) +} + +func TestTopologicalOrderCycleDetection(t *testing.T) { + states := map[string]*nodeState{ + "a": newNodeState("a", []string{"c"}), + "b": newNodeState("b", []string{"a"}), + "c": newNodeState("c", []string{"b"}), + } + + _, err := topologicalOrder(states) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cycle detected") +} + +func TestTopologicalOrderSingleNode(t *testing.T) { + states := map[string]*nodeState{ + "only": newNodeState("only", nil), + } + + waves, err := topologicalOrder(states) + require.NoError(t, err) + require.Len(t, waves, 1) + assert.Equal(t, []string{"only"}, waves[0]) +} + +func TestResolveRefs(t *testing.T) { + states := map[string]*nodeState{ + "search": newNodeState("search", nil), + } + states["search"].setResult("found: /tmp/file.txt", nil) + + input := `{"path":"#nodesearch"}` + result := resolveRefs(input, states) + assert.Contains(t, result, "found: /tmp/file.txt") + assert.NotContains(t, result, "#nodesearch") +} + +func TestResolveRefsNoMatch(t *testing.T) { + states := map[string]*nodeState{} + input := `{"path":"#nodemissing"}` + result := resolveRefs(input, states) + assert.Equal(t, input, result) +} + +func TestResolveToolExecArgsNoRefs(t *testing.T) { + states := map[string]*nodeState{} + input := `{"path":"/tmp/plain.txt"}` + result := resolveToolExecArgs(input, states) + assert.Equal(t, input, result) +} + +func TestEscapeForJSON(t *testing.T) { + tests := []struct { + input string + expected string + }{ + {"hello", "hello"}, + {`has "quotes"`, `has \"quotes\"`}, + {"has\nnewline", `has\nnewline`}, + {"has\ttab", `has\ttab`}, + } + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + assert.Equal(t, tt.expected, escapeForJSON(tt.input)) + }) + } +} + +func TestNodeStateSetAndGetResult(t *testing.T) { + ns := newNodeState("test", nil) + + go func() { + ns.setResult("result-data", nil) + }() + + <-ns.done + + result, err := ns.getResult() + assert.NoError(t, err) + assert.Equal(t, "result-data", result) +} diff --git a/pkg/itr/dag/router_test.go b/pkg/itr/dag/router_test.go new file mode 100644 index 000000000..c66d08abc --- /dev/null +++ b/pkg/itr/dag/router_test.go @@ -0,0 +1,60 @@ +package dag + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRouteExplicitModes(t *testing.T) { + cfg := DefaultRouterConfig() + + assert.Equal(t, ModeReAct, Route(ModeReAct, "anything", cfg)) + assert.Equal(t, ModeDAG, Route(ModeDAG, "anything", cfg)) +} + +func TestRouteAutoSimpleQuery(t *testing.T) { + cfg := DefaultRouterConfig() + assert.Equal(t, ModeReAct, Route(ModeAuto, "what is the weather?", cfg)) +} + +func TestRouteAutoComplexQuery(t *testing.T) { + cfg := DefaultRouterConfig() + longQuery := strings.Repeat("word ", 35) + assert.Equal(t, ModeDAG, Route(ModeAuto, longQuery, cfg)) +} + +func TestRouteAutoParallelKeywords(t *testing.T) { + cfg := DefaultRouterConfig() + + keywords := []string{ + "search files and then fetch URL", + "do both tasks simultaneously", + "compare results from multiple sources", + "aggregate data from several endpoints", + } + for _, q := range keywords { + t.Run(q, func(t *testing.T) { + assert.Equal(t, ModeDAG, Route(ModeAuto, q, cfg)) + }) + } +} + +func TestRouteAutoToolSignals(t *testing.T) { + cfg := DefaultRouterConfig() + q := "search the codebase, read the file, then execute the command" + assert.Equal(t, ModeDAG, Route(ModeAuto, q, cfg)) +} + +func TestToolLoopModeString(t *testing.T) { + assert.Equal(t, "react", ModeReAct.String()) + assert.Equal(t, "dag", ModeDAG.String()) + assert.Equal(t, "auto", ModeAuto.String()) + assert.Equal(t, "unknown", ToolLoopMode(99).String()) +} + +func TestClassifyQueryDefault(t *testing.T) { + cfg := DefaultRouterConfig() + assert.Equal(t, ModeReAct, classifyQuery("hello", cfg)) +} diff --git a/pkg/rlm/fanout_test.go b/pkg/rlm/fanout_test.go new file mode 100644 index 000000000..f42d8d472 --- /dev/null +++ b/pkg/rlm/fanout_test.go @@ -0,0 +1,130 @@ +package rlm + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestFanOutEmpty(t *testing.T) { + results := FanOut(context.Background(), nil, 0, nil) + assert.Nil(t, results) +} + +func TestFanOutUnbounded(t *testing.T) { + partitions := []string{"part-0", "part-1", "part-2"} + + results := FanOut(context.Background(), partitions, 0, func(ctx context.Context, idx int, key, partition string) PartitionResult { + return PartitionResult{ + PartitionIdx: idx, + ContextKey: key, + Answer: fmt.Sprintf("answer for %s", partition), + Tokens: 10, + } + }) + + require.Len(t, results, 3) + for i, r := range results { + assert.Equal(t, i, r.PartitionIdx) + assert.Contains(t, r.Answer, fmt.Sprintf("part-%d", i)) + assert.Equal(t, uint32(10), r.Tokens) + } +} + +func TestFanOutBounded(t *testing.T) { + partitions := make([]string, 10) + for i := range partitions { + partitions[i] = fmt.Sprintf("chunk-%d", i) + } + + var maxConcurrent int64 + var current int64 + + results := FanOut(context.Background(), partitions, 3, func(ctx context.Context, idx int, key, partition string) PartitionResult { + c := atomic.AddInt64(¤t, 1) + for { + old := atomic.LoadInt64(&maxConcurrent) + if c <= old { + break + } + if atomic.CompareAndSwapInt64(&maxConcurrent, old, c) { + break + } + } + time.Sleep(5 * time.Millisecond) + atomic.AddInt64(¤t, -1) + return PartitionResult{ + PartitionIdx: idx, + Answer: partition, + Tokens: 1, + } + }) + + require.Len(t, results, 10) + assert.LessOrEqual(t, atomic.LoadInt64(&maxConcurrent), int64(3), + "max concurrent goroutines should respect the limit") +} + +func TestFanOutPreservesOrder(t *testing.T) { + partitions := []string{"A", "B", "C", "D"} + + results := FanOut(context.Background(), partitions, 2, func(ctx context.Context, idx int, key, partition string) PartitionResult { + return PartitionResult{ + PartitionIdx: idx, + Answer: partition, + } + }) + + require.Len(t, results, 4) + for i, r := range results { + assert.Equal(t, i, r.PartitionIdx) + assert.Equal(t, partitions[i], r.Answer) + } +} + +func TestMergeResultsDeduplication(t *testing.T) { + results := []PartitionResult{ + {Answer: " answer one "}, + {Answer: "answer one"}, + {Answer: "answer two"}, + {Answer: "", Err: fmt.Errorf("failed")}, + {Answer: "answer two"}, + } + + merged := MergeResults(results) + assert.Equal(t, "answer one\nanswer two", merged) +} + +func TestMergeResultsAllErrors(t *testing.T) { + results := []PartitionResult{ + {Err: fmt.Errorf("e1")}, + {Err: fmt.Errorf("e2")}, + } + assert.Empty(t, MergeResults(results)) +} + +func TestMergeResultsAllEmpty(t *testing.T) { + results := []PartitionResult{ + {Answer: ""}, + {Answer: " "}, + } + assert.Empty(t, MergeResults(results)) +} + +func TestTotalTokens(t *testing.T) { + results := []PartitionResult{ + {Tokens: 100}, + {Tokens: 250}, + {Tokens: 50}, + } + assert.Equal(t, uint32(400), TotalTokens(results)) +} + +func TestTotalTokensEmpty(t *testing.T) { + assert.Equal(t, uint32(0), TotalTokens(nil)) +} diff --git a/pkg/rlm/strategy_test.go b/pkg/rlm/strategy_test.go new file mode 100644 index 000000000..3453233b3 --- /dev/null +++ b/pkg/rlm/strategy_test.go @@ -0,0 +1,86 @@ +package rlm + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestStrategyPlanNextFinalAtMaxDepth(t *testing.T) { + sp := NewStrategyPlanner(StrategyConfig{ + DirectThreshold: 8192, + DefaultPartitionK: 4, + MaxDepth: 3, + }) + + op := sp.PlanNext(100000, "any query", 3) + assert.Equal(t, OpFinal, op.Type) +} + +func TestStrategyPlanNextFinalSmallContext(t *testing.T) { + sp := NewStrategyPlanner(DefaultStrategyConfig()) + + op := sp.PlanNext(1000, "any query", 0) + assert.Equal(t, OpFinal, op.Type) +} + +func TestStrategyPlanNextGrepForKeywordQuery(t *testing.T) { + sp := NewStrategyPlanner(DefaultStrategyConfig()) + + tests := []struct { + query string + }{ + {`find "handleRequest"`}, + {"search for main function"}, + {"where is the config loaded"}, + {"error: connection refused trace"}, + {"func processData handler"}, + {"def calculate_total in utils"}, + } + for _, tt := range tests { + t.Run(tt.query, func(t *testing.T) { + op := sp.PlanNext(1_000_000, tt.query, 0) + assert.Equal(t, OpGrep, op.Type) + assert.NotEmpty(t, op.GrepQuery) + }) + } +} + +func TestStrategyPlanNextPartitionDefault(t *testing.T) { + sp := NewStrategyPlanner(DefaultStrategyConfig()) + op := sp.PlanNext(100_000, "summarize this document", 0) + assert.Equal(t, OpPartition, op.Type) + assert.Equal(t, 4, op.PartitionK) +} + +func TestStrategyPlanNextPartitionLargeContext(t *testing.T) { + sp := NewStrategyPlanner(DefaultStrategyConfig()) + op := sp.PlanNext(5_000_000, "summarize this corpus", 0) + assert.Equal(t, OpPartition, op.Type) + assert.Equal(t, 8, op.PartitionK, "large contexts should use more partitions") +} + +func TestExtractKeywordQuoted(t *testing.T) { + assert.Equal(t, "handleRequest", extractKeyword(`find "handleRequest" in the codebase`)) +} + +func TestExtractKeywordNoQuotes(t *testing.T) { + assert.Equal(t, "find", extractKeyword("find the main function")) +} + +func TestExtractKeywordEmpty(t *testing.T) { + assert.Equal(t, "", extractKeyword("")) +} + +func TestLooksLikeKeywordQuery(t *testing.T) { + assert.True(t, looksLikeKeywordQuery(`find "something"`)) + assert.True(t, looksLikeKeywordQuery("error: something broke")) + assert.True(t, looksLikeKeywordQuery("func processData")) + assert.True(t, looksLikeKeywordQuery("def main")) + assert.True(t, looksLikeKeywordQuery("find the bug")) + assert.True(t, looksLikeKeywordQuery("search for pattern")) + assert.True(t, looksLikeKeywordQuery("where is the config")) + + assert.False(t, looksLikeKeywordQuery("summarize this")) + assert.False(t, looksLikeKeywordQuery("how does X work")) +} diff --git a/pkg/security/securebus/audit_test.go b/pkg/security/securebus/audit_test.go new file mode 100644 index 000000000..8d1dab270 --- /dev/null +++ b/pkg/security/securebus/audit_test.go @@ -0,0 +1,132 @@ +package securebus + +import ( + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAuditLogAppendAndRetrieve(t *testing.T) { + al := NewAuditLog() + + event := AuditEvent{ + RequestID: "req-1", + SessionKey: "sess-A", + ToolName: "read_file", + CommandType: "tool_exec", + At: time.Now(), + DurationMS: 50, + } + + require.NoError(t, al.Append(event)) + assert.Equal(t, 1, al.Len()) + + events := al.Events() + require.Len(t, events, 1) + assert.Equal(t, "req-1", events[0].RequestID) + assert.Equal(t, "read_file", events[0].ToolName) +} + +func TestAuditLogConcurrentAppend(t *testing.T) { + al := NewAuditLog() + n := 100 + + var wg sync.WaitGroup + wg.Add(n) + for i := 0; i < n; i++ { + go func(idx int) { + defer wg.Done() + _ = al.Append(AuditEvent{ + RequestID: fmt.Sprintf("req-%d", idx), + At: time.Now(), + }) + }(i) + } + wg.Wait() + + assert.Equal(t, n, al.Len()) +} + +func TestAuditLogFilterBySession(t *testing.T) { + al := NewAuditLog() + + _ = al.Append(AuditEvent{RequestID: "r1", SessionKey: "sess-A"}) + _ = al.Append(AuditEvent{RequestID: "r2", SessionKey: "sess-B"}) + _ = al.Append(AuditEvent{RequestID: "r3", SessionKey: "sess-A"}) + + sessA := al.FilterBySession("sess-A") + assert.Len(t, sessA, 2) + + sessB := al.FilterBySession("sess-B") + assert.Len(t, sessB, 1) + + sessC := al.FilterBySession("sess-C") + assert.Empty(t, sessC) +} + +func TestAuditLogLeakEvents(t *testing.T) { + al := NewAuditLog() + + _ = al.Append(AuditEvent{RequestID: "r1", LeakDetected: false}) + _ = al.Append(AuditEvent{RequestID: "r2", LeakDetected: true, RedactedKeys: []string{"api_key"}}) + _ = al.Append(AuditEvent{RequestID: "r3", LeakDetected: true, RedactedKeys: []string{"token"}}) + + leaks := al.LeakEvents() + assert.Len(t, leaks, 2) + assert.Equal(t, "r2", leaks[0].RequestID) + assert.Equal(t, "r3", leaks[1].RequestID) +} + +type mockSink struct { + mu sync.Mutex + events []AuditEvent + failAt int +} + +func (ms *mockSink) Write(event AuditEvent) error { + ms.mu.Lock() + defer ms.mu.Unlock() + if ms.failAt > 0 && len(ms.events) >= ms.failAt { + return fmt.Errorf("sink full") + } + ms.events = append(ms.events, event) + return nil +} + +func TestAuditLogSinkIntegration(t *testing.T) { + sink := &mockSink{} + al := NewAuditLog(sink) + + _ = al.Append(AuditEvent{RequestID: "r1"}) + _ = al.Append(AuditEvent{RequestID: "r2"}) + + assert.Equal(t, 2, al.Len()) + assert.Len(t, sink.events, 2) +} + +func TestAuditLogSinkError(t *testing.T) { + sink := &mockSink{failAt: 1} + al := NewAuditLog(sink) + + assert.NoError(t, al.Append(AuditEvent{RequestID: "r1"})) + + err := al.Append(AuditEvent{RequestID: "r2"}) + assert.Error(t, err) + + assert.Equal(t, 2, al.Len(), "in-memory log should always append") +} + +func TestAuditLogEventsImmutable(t *testing.T) { + al := NewAuditLog() + _ = al.Append(AuditEvent{RequestID: "r1"}) + + events := al.Events() + events[0].RequestID = "mutated" + + original := al.Events() + assert.Equal(t, "r1", original[0].RequestID, "original should be unaffected") +} diff --git a/pkg/security/securebus/policy_test.go b/pkg/security/securebus/policy_test.go new file mode 100644 index 000000000..3f5d57f88 --- /dev/null +++ b/pkg/security/securebus/policy_test.go @@ -0,0 +1,103 @@ +package securebus + +import ( + "testing" + + "github.com/sipeed/picoclaw/pkg/itr" + "github.com/sipeed/picoclaw/pkg/tools" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPolicyValidateRecursionDepth(t *testing.T) { + pe := NewPolicyEngine(PolicyConfig{MaxRecursionDepth: 5}) + + req := itr.ToolRequest{Depth: 3} + assert.NoError(t, pe.Validate(req, tools.ZeroCapabilities())) + + req.Depth = 6 + err := pe.Validate(req, tools.ZeroCapabilities()) + require.Error(t, err) + assert.Contains(t, err.Error(), "recursion depth") +} + +func TestPolicyValidateNoDepthLimit(t *testing.T) { + pe := NewPolicyEngine(PolicyConfig{MaxRecursionDepth: 0}) + + req := itr.ToolRequest{Depth: 255} + assert.NoError(t, pe.Validate(req, tools.ZeroCapabilities())) +} + +func TestPolicyValidateNetworkSSRFBlocked(t *testing.T) { + pe := NewPolicyEngine(DefaultPolicyConfig()) + + ssrfURLs := []string{ + "http://169.254.169.254/latest/meta-data/", + "http://10.0.0.1/internal", + "http://172.16.0.1/private", + "http://192.168.1.1/admin", + "http://localhost:8080/health", + "http://127.0.0.1:3000/api", + "https://169.254.169.254/token", + } + + rules := []tools.EndpointRule{{Pattern: "*"}} + + for _, url := range ssrfURLs { + t.Run(url, func(t *testing.T) { + err := pe.ValidateNetwork(url, rules) + require.Error(t, err) + assert.Contains(t, err.Error(), "SSRF blocklist") + }) + } +} + +func TestPolicyValidateNetworkAllowed(t *testing.T) { + pe := NewPolicyEngine(DefaultPolicyConfig()) + rules := []tools.EndpointRule{{Pattern: "https://api.github.com/*"}} + + assert.NoError(t, pe.ValidateNetwork("https://api.github.com/repos", rules)) +} + +func TestPolicyValidateNetworkNoRules(t *testing.T) { + pe := NewPolicyEngine(DefaultPolicyConfig()) + err := pe.ValidateNetwork("https://example.com", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "no EndpointRules") +} + +func TestPolicyValidateNetworkNoMatchingRule(t *testing.T) { + pe := NewPolicyEngine(DefaultPolicyConfig()) + rules := []tools.EndpointRule{{Pattern: "https://api.github.com/*"}} + + err := pe.ValidateNetwork("https://evil.com/steal", rules) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not match") +} + +func TestPolicyValidateFilesystemAllowed(t *testing.T) { + pe := NewPolicyEngine(PolicyConfig{AllowedWorkspace: "/workspace"}) + rules := []tools.PathRule{{Pattern: "src/*", Mode: "rw"}} + + assert.NoError(t, pe.ValidateFilesystem("/workspace/src/main.go", "r", rules)) + assert.NoError(t, pe.ValidateFilesystem("/workspace/src/main.go", "w", rules)) + assert.NoError(t, pe.ValidateFilesystem("/workspace/src/main.go", "rw", rules)) +} + +func TestPolicyValidateFilesystemNoRules(t *testing.T) { + pe := NewPolicyEngine(PolicyConfig{}) + err := pe.ValidateFilesystem("/etc/passwd", "r", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "no PathRules") +} + +func TestPolicyValidateFilesystemModeMismatch(t *testing.T) { + pe := NewPolicyEngine(PolicyConfig{AllowedWorkspace: "/workspace"}) + rules := []tools.PathRule{{Pattern: "data/*", Mode: "r"}} + + assert.NoError(t, pe.ValidateFilesystem("/workspace/data/file.csv", "r", rules)) + + err := pe.ValidateFilesystem("/workspace/data/file.csv", "w", rules) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not match") +}