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
This commit is contained in:
ZanzyTHEbar 2026-02-19 14:48:46 +00:00
parent 364746307f
commit 1f74b227b9
9 changed files with 1061 additions and 0 deletions

228
pkg/itr/commands_test.go Normal file
View file

@ -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)
}

184
pkg/itr/dag/planner_test.go Normal file
View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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))
}

130
pkg/rlm/fanout_test.go Normal file
View file

@ -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(&current, 1)
for {
old := atomic.LoadInt64(&maxConcurrent)
if c <= old {
break
}
if atomic.CompareAndSwapInt64(&maxConcurrent, old, c) {
break
}
}
time.Sleep(5 * time.Millisecond)
atomic.AddInt64(&current, -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))
}

86
pkg/rlm/strategy_test.go Normal file
View file

@ -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"))
}

View file

@ -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")
}

View file

@ -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")
}