Merge pull request #13 from HarnessLab/copilot/ensure-features-are-presented
Add dedicated test coverage for all untested feature modules
This commit is contained in:
Abdelrahman Abdallah
+2
-2
@@ -910,8 +910,8 @@ def _ensure_shell_allowed(command: str, context: ToolExecutionContext) -> None:
|
||||
r'(^|[;&|])\s*shutdown\s',
|
||||
r'(^|[;&|])\s*reboot\s',
|
||||
r'(^|[;&|])\s*mkfs',
|
||||
r'(^|[;&|])\s*chmod\s+-R\s+777',
|
||||
r'(^|[;&|])\s*chown\s+-R',
|
||||
r'(^|[;&|])\s*chmod\s+-r\s+777',
|
||||
r'(^|[;&|])\s*chown\s+-r',
|
||||
r'(^|[;&|])\s*git\s+reset\s+--hard',
|
||||
r'(^|[;&|])\s*git\s+clean\s+-fd',
|
||||
r'(^|[;&|])\s*:\s*>\s*',
|
||||
|
||||
@@ -0,0 +1,581 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.agent_manager import AgentManager, ManagedAgentGroup, ManagedAgentRecord
|
||||
|
||||
|
||||
class TestManagedAgentRecordDefaults(unittest.TestCase):
|
||||
def test_defaults(self) -> None:
|
||||
rec = ManagedAgentRecord(agent_id="a1", prompt="do stuff")
|
||||
self.assertEqual(rec.agent_id, "a1")
|
||||
self.assertEqual(rec.prompt, "do stuff")
|
||||
self.assertIsNone(rec.parent_agent_id)
|
||||
self.assertIsNone(rec.group_id)
|
||||
self.assertIsNone(rec.child_index)
|
||||
self.assertIsNone(rec.label)
|
||||
self.assertIsNone(rec.resumed_from_session_id)
|
||||
self.assertIsNone(rec.session_id)
|
||||
self.assertIsNone(rec.session_path)
|
||||
self.assertEqual(rec.status, "running")
|
||||
self.assertEqual(rec.turns, 0)
|
||||
self.assertEqual(rec.tool_calls, 0)
|
||||
self.assertIsNone(rec.stop_reason)
|
||||
|
||||
|
||||
class TestManagedAgentGroupDefaults(unittest.TestCase):
|
||||
def test_defaults(self) -> None:
|
||||
grp = ManagedAgentGroup(group_id="g1")
|
||||
self.assertEqual(grp.group_id, "g1")
|
||||
self.assertIsNone(grp.label)
|
||||
self.assertIsNone(grp.parent_agent_id)
|
||||
self.assertEqual(grp.child_agent_ids, ())
|
||||
self.assertEqual(grp.strategy, "serial")
|
||||
self.assertEqual(grp.status, "running")
|
||||
self.assertEqual(grp.completed_children, 0)
|
||||
self.assertEqual(grp.failed_children, 0)
|
||||
self.assertEqual(grp.batch_count, 0)
|
||||
self.assertEqual(grp.max_batch_size, 0)
|
||||
self.assertEqual(grp.dependency_skips, 0)
|
||||
|
||||
|
||||
class TestStartAgent(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_increments_counter_and_returns_unique_ids(self) -> None:
|
||||
id1 = self.mgr.start_agent(prompt="task1")
|
||||
id2 = self.mgr.start_agent(prompt="task2")
|
||||
id3 = self.mgr.start_agent(prompt="task3")
|
||||
self.assertEqual(id1, "agent_1")
|
||||
self.assertEqual(id2, "agent_2")
|
||||
self.assertEqual(id3, "agent_3")
|
||||
self.assertEqual(len(self.mgr.records), 3)
|
||||
|
||||
def test_record_stored_with_correct_fields(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="hello", label="my-label")
|
||||
rec = self.mgr.records[aid]
|
||||
self.assertEqual(rec.agent_id, aid)
|
||||
self.assertEqual(rec.prompt, "hello")
|
||||
self.assertEqual(rec.label, "my-label")
|
||||
self.assertEqual(rec.status, "running")
|
||||
|
||||
def test_with_parent_agent_id_tracks_lineage(self) -> None:
|
||||
parent = self.mgr.start_agent(prompt="parent")
|
||||
child = self.mgr.start_agent(prompt="child", parent_agent_id=parent)
|
||||
rec = self.mgr.records[child]
|
||||
self.assertEqual(rec.parent_agent_id, parent)
|
||||
|
||||
def test_with_group_id_registers_child(self) -> None:
|
||||
gid = self.mgr.start_group(label="grp")
|
||||
aid = self.mgr.start_agent(prompt="task", group_id=gid, child_index=0)
|
||||
grp = self.mgr.groups[gid]
|
||||
self.assertIn(aid, grp.child_agent_ids)
|
||||
self.assertEqual(self.mgr.records[aid].group_id, gid)
|
||||
self.assertEqual(self.mgr.records[aid].child_index, 0)
|
||||
|
||||
def test_resumed_agents_tracked(self) -> None:
|
||||
aid = self.mgr.start_agent(
|
||||
prompt="resume", resumed_from_session_id="sess-old-123"
|
||||
)
|
||||
rec = self.mgr.records[aid]
|
||||
self.assertEqual(rec.resumed_from_session_id, "sess-old-123")
|
||||
|
||||
|
||||
class TestStartGroup(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_creates_group_with_strategy(self) -> None:
|
||||
gid = self.mgr.start_group(label="batch", strategy="parallel")
|
||||
self.assertEqual(gid, "group_1")
|
||||
grp = self.mgr.groups[gid]
|
||||
self.assertEqual(grp.label, "batch")
|
||||
self.assertEqual(grp.strategy, "parallel")
|
||||
self.assertEqual(grp.status, "running")
|
||||
|
||||
def test_increments_group_counter(self) -> None:
|
||||
g1 = self.mgr.start_group(label="a")
|
||||
g2 = self.mgr.start_group(label="b")
|
||||
self.assertEqual(g1, "group_1")
|
||||
self.assertEqual(g2, "group_2")
|
||||
|
||||
def test_parent_agent_id_stored(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="parent")
|
||||
gid = self.mgr.start_group(label="child-group", parent_agent_id=aid)
|
||||
self.assertEqual(self.mgr.groups[gid].parent_agent_id, aid)
|
||||
|
||||
def test_default_strategy_is_serial(self) -> None:
|
||||
gid = self.mgr.start_group()
|
||||
self.assertEqual(self.mgr.groups[gid].strategy, "serial")
|
||||
|
||||
|
||||
class TestRegisterGroupChild(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_adds_agent_to_group(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
aid = self.mgr.start_agent(prompt="t")
|
||||
self.mgr.register_group_child(gid, aid, child_index=0)
|
||||
grp = self.mgr.groups[gid]
|
||||
self.assertIn(aid, grp.child_agent_ids)
|
||||
self.assertEqual(self.mgr.records[aid].group_id, gid)
|
||||
self.assertEqual(self.mgr.records[aid].child_index, 0)
|
||||
|
||||
def test_duplicate_does_not_add_twice(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
aid = self.mgr.start_agent(prompt="t", group_id=gid, child_index=0)
|
||||
# Already registered via start_agent; register again
|
||||
self.mgr.register_group_child(gid, aid, child_index=0)
|
||||
grp = self.mgr.groups[gid]
|
||||
self.assertEqual(grp.child_agent_ids.count(aid), 1)
|
||||
|
||||
def test_unknown_group_is_noop(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="t")
|
||||
# Should not raise
|
||||
self.mgr.register_group_child("nonexistent", aid, child_index=0)
|
||||
|
||||
def test_unknown_agent_does_not_crash(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
# Agent does not exist; group gets the ID but record update is skipped
|
||||
self.mgr.register_group_child(gid, "fake_agent", child_index=0)
|
||||
self.assertIn("fake_agent", self.mgr.groups[gid].child_agent_ids)
|
||||
|
||||
def test_updates_child_index_on_record(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
aid = self.mgr.start_agent(prompt="t")
|
||||
self.mgr.register_group_child(gid, aid, child_index=5)
|
||||
self.assertEqual(self.mgr.records[aid].child_index, 5)
|
||||
self.assertEqual(self.mgr.records[aid].group_id, gid)
|
||||
|
||||
|
||||
class TestFinishAgent(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_marks_completed_with_stats(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="work")
|
||||
self.mgr.finish_agent(
|
||||
aid,
|
||||
session_id="sess-1",
|
||||
session_path="/path/sess",
|
||||
turns=5,
|
||||
tool_calls=12,
|
||||
stop_reason="end_turn",
|
||||
)
|
||||
rec = self.mgr.records[aid]
|
||||
self.assertEqual(rec.status, "completed")
|
||||
self.assertEqual(rec.session_id, "sess-1")
|
||||
self.assertEqual(rec.session_path, "/path/sess")
|
||||
self.assertEqual(rec.turns, 5)
|
||||
self.assertEqual(rec.tool_calls, 12)
|
||||
self.assertEqual(rec.stop_reason, "end_turn")
|
||||
|
||||
def test_preserves_original_fields(self) -> None:
|
||||
aid = self.mgr.start_agent(
|
||||
prompt="p", parent_agent_id="parent_x", label="lbl"
|
||||
)
|
||||
self.mgr.finish_agent(
|
||||
aid,
|
||||
session_id="s",
|
||||
session_path="/p",
|
||||
turns=1,
|
||||
tool_calls=2,
|
||||
stop_reason=None,
|
||||
)
|
||||
rec = self.mgr.records[aid]
|
||||
self.assertEqual(rec.prompt, "p")
|
||||
self.assertEqual(rec.parent_agent_id, "parent_x")
|
||||
self.assertEqual(rec.label, "lbl")
|
||||
|
||||
def test_unknown_agent_is_noop(self) -> None:
|
||||
# Should not raise
|
||||
self.mgr.finish_agent(
|
||||
"unknown_id",
|
||||
session_id=None,
|
||||
session_path=None,
|
||||
turns=0,
|
||||
tool_calls=0,
|
||||
stop_reason=None,
|
||||
)
|
||||
self.assertEqual(len(self.mgr.records), 0)
|
||||
|
||||
|
||||
class TestFinishGroup(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_updates_group_status_and_stats(self) -> None:
|
||||
gid = self.mgr.start_group(label="g", strategy="parallel")
|
||||
self.mgr.finish_group(
|
||||
gid,
|
||||
status="completed",
|
||||
completed_children=3,
|
||||
failed_children=1,
|
||||
batch_count=2,
|
||||
max_batch_size=4,
|
||||
dependency_skips=0,
|
||||
)
|
||||
grp = self.mgr.groups[gid]
|
||||
self.assertEqual(grp.status, "completed")
|
||||
self.assertEqual(grp.completed_children, 3)
|
||||
self.assertEqual(grp.failed_children, 1)
|
||||
self.assertEqual(grp.batch_count, 2)
|
||||
self.assertEqual(grp.max_batch_size, 4)
|
||||
self.assertEqual(grp.dependency_skips, 0)
|
||||
# Preserved fields
|
||||
self.assertEqual(grp.label, "g")
|
||||
self.assertEqual(grp.strategy, "parallel")
|
||||
|
||||
def test_unknown_group_is_noop(self) -> None:
|
||||
self.mgr.finish_group(
|
||||
"ghost",
|
||||
status="completed",
|
||||
completed_children=0,
|
||||
failed_children=0,
|
||||
)
|
||||
self.assertEqual(len(self.mgr.groups), 0)
|
||||
|
||||
def test_preserves_child_agent_ids(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
aid = self.mgr.start_agent(prompt="t", group_id=gid, child_index=0)
|
||||
self.mgr.finish_group(
|
||||
gid, status="completed", completed_children=1, failed_children=0
|
||||
)
|
||||
self.assertIn(aid, self.mgr.groups[gid].child_agent_ids)
|
||||
|
||||
|
||||
class TestChildrenOf(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_returns_only_children_of_specified_parent(self) -> None:
|
||||
p1 = self.mgr.start_agent(prompt="parent1")
|
||||
p2 = self.mgr.start_agent(prompt="parent2")
|
||||
c1 = self.mgr.start_agent(prompt="c1", parent_agent_id=p1)
|
||||
c2 = self.mgr.start_agent(prompt="c2", parent_agent_id=p1)
|
||||
c3 = self.mgr.start_agent(prompt="c3", parent_agent_id=p2)
|
||||
|
||||
children_p1 = self.mgr.children_of(p1)
|
||||
children_p2 = self.mgr.children_of(p2)
|
||||
|
||||
self.assertEqual(len(children_p1), 2)
|
||||
ids_p1 = {r.agent_id for r in children_p1}
|
||||
self.assertEqual(ids_p1, {c1, c2})
|
||||
|
||||
self.assertEqual(len(children_p2), 1)
|
||||
self.assertEqual(children_p2[0].agent_id, c3)
|
||||
|
||||
def test_returns_empty_for_no_children(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="solo")
|
||||
self.assertEqual(self.mgr.children_of(aid), ())
|
||||
|
||||
def test_returns_empty_for_unknown_parent(self) -> None:
|
||||
self.assertEqual(self.mgr.children_of("nonexistent"), ())
|
||||
|
||||
|
||||
class TestGroupChildren(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_returns_sorted_members(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
a2 = self.mgr.start_agent(prompt="b", group_id=gid, child_index=2)
|
||||
a0 = self.mgr.start_agent(prompt="a", group_id=gid, child_index=0)
|
||||
a1 = self.mgr.start_agent(prompt="c", group_id=gid, child_index=1)
|
||||
|
||||
children = self.mgr.group_children(gid)
|
||||
self.assertEqual(len(children), 3)
|
||||
self.assertEqual(children[0].agent_id, a0)
|
||||
self.assertEqual(children[1].agent_id, a1)
|
||||
self.assertEqual(children[2].agent_id, a2)
|
||||
|
||||
def test_none_child_index_sorted_last(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
a_none = self.mgr.start_agent(prompt="x", group_id=gid)
|
||||
a0 = self.mgr.start_agent(prompt="y", group_id=gid, child_index=0)
|
||||
|
||||
children = self.mgr.group_children(gid)
|
||||
self.assertEqual(children[0].agent_id, a0)
|
||||
self.assertEqual(children[1].agent_id, a_none)
|
||||
|
||||
def test_empty_for_unknown_group(self) -> None:
|
||||
self.assertEqual(self.mgr.group_children("nope"), ())
|
||||
|
||||
|
||||
class TestGroupSummary(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_aggregates_statistics(self) -> None:
|
||||
gid = self.mgr.start_group(label="batch", strategy="parallel")
|
||||
a1 = self.mgr.start_agent(
|
||||
prompt="t1", group_id=gid, child_index=0,
|
||||
resumed_from_session_id="old-sess",
|
||||
)
|
||||
a2 = self.mgr.start_agent(prompt="t2", group_id=gid, child_index=1)
|
||||
self.mgr.finish_agent(
|
||||
a1, session_id="s1", session_path="/p1",
|
||||
turns=3, tool_calls=5, stop_reason="end_turn",
|
||||
)
|
||||
self.mgr.finish_agent(
|
||||
a2, session_id="s2", session_path="/p2",
|
||||
turns=2, tool_calls=4, stop_reason="max_turns",
|
||||
)
|
||||
self.mgr.finish_group(
|
||||
gid, status="completed",
|
||||
completed_children=2, failed_children=0,
|
||||
batch_count=1, max_batch_size=2,
|
||||
)
|
||||
|
||||
summary = self.mgr.group_summary(gid)
|
||||
assert summary is not None
|
||||
self.assertEqual(summary["group_id"], gid)
|
||||
self.assertEqual(summary["label"], "batch")
|
||||
self.assertEqual(summary["strategy"], "parallel")
|
||||
self.assertEqual(summary["status"], "completed")
|
||||
self.assertEqual(summary["child_count"], 2)
|
||||
self.assertEqual(summary["completed_children"], 2)
|
||||
self.assertEqual(summary["failed_children"], 0)
|
||||
self.assertEqual(summary["resumed_children"], 1)
|
||||
self.assertEqual(summary["batch_count"], 1)
|
||||
self.assertEqual(summary["max_batch_size"], 2)
|
||||
self.assertEqual(summary["dependency_skips"], 0)
|
||||
self.assertEqual(
|
||||
summary["stop_reason_counts"],
|
||||
{"end_turn": 1, "max_turns": 1},
|
||||
)
|
||||
|
||||
def test_running_agents_counted_as_na(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
self.mgr.start_agent(prompt="t", group_id=gid, child_index=0)
|
||||
summary = self.mgr.group_summary(gid)
|
||||
assert summary is not None
|
||||
self.assertEqual(summary["stop_reason_counts"], {"n/a": 1})
|
||||
|
||||
def test_returns_none_for_unknown_group(self) -> None:
|
||||
self.assertIsNone(self.mgr.group_summary("unknown"))
|
||||
|
||||
|
||||
class TestCompletedRecords(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_filters_only_completed(self) -> None:
|
||||
a1 = self.mgr.start_agent(prompt="t1")
|
||||
a2 = self.mgr.start_agent(prompt="t2")
|
||||
a3 = self.mgr.start_agent(prompt="t3")
|
||||
self.mgr.finish_agent(
|
||||
a1, session_id="s", session_path="/p",
|
||||
turns=1, tool_calls=1, stop_reason="done",
|
||||
)
|
||||
self.mgr.finish_agent(
|
||||
a3, session_id="s2", session_path="/p2",
|
||||
turns=2, tool_calls=3, stop_reason="done",
|
||||
)
|
||||
|
||||
completed = self.mgr.completed_records()
|
||||
self.assertEqual(len(completed), 2)
|
||||
ids = {r.agent_id for r in completed}
|
||||
self.assertEqual(ids, {a1, a3})
|
||||
|
||||
def test_empty_when_none_completed(self) -> None:
|
||||
self.mgr.start_agent(prompt="running")
|
||||
self.assertEqual(self.mgr.completed_records(), ())
|
||||
|
||||
|
||||
class TestSummaryLines(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.mgr = AgentManager()
|
||||
|
||||
def test_empty_manager(self) -> None:
|
||||
lines = self.mgr.summary_lines()
|
||||
self.assertIn("- Managed agents: 0", lines)
|
||||
self.assertIn("- Completed agents: 0", lines)
|
||||
self.assertIn("- Child agents: 0", lines)
|
||||
self.assertIn("- Resumed agents: 0", lines)
|
||||
self.assertIn("- Agent groups: 0", lines)
|
||||
self.assertIn("- Completed groups: 0", lines)
|
||||
|
||||
def test_basic_output_format(self) -> None:
|
||||
a1 = self.mgr.start_agent(prompt="task", label="worker-1")
|
||||
self.mgr.finish_agent(
|
||||
a1, session_id="s1", session_path="/p",
|
||||
turns=4, tool_calls=10, stop_reason="end_turn",
|
||||
)
|
||||
lines = self.mgr.summary_lines()
|
||||
self.assertIn("- Managed agents: 1", lines)
|
||||
self.assertIn("- Completed agents: 1", lines)
|
||||
# Agent detail line
|
||||
detail = [l for l in lines if "worker-1" in l]
|
||||
self.assertEqual(len(detail), 1)
|
||||
self.assertIn("status=completed", detail[0])
|
||||
self.assertIn("turns=4", detail[0])
|
||||
self.assertIn("tool_calls=10", detail[0])
|
||||
self.assertIn("stop=end_turn", detail[0])
|
||||
|
||||
def test_group_info_in_agent_line(self) -> None:
|
||||
gid = self.mgr.start_group(label="g")
|
||||
self.mgr.start_agent(prompt="t", group_id=gid, child_index=0, label="child-0")
|
||||
lines = self.mgr.summary_lines()
|
||||
detail = [l for l in lines if "child-0" in l]
|
||||
self.assertEqual(len(detail), 1)
|
||||
self.assertIn(f"group={gid}", detail[0])
|
||||
self.assertIn("child_index=0", detail[0])
|
||||
|
||||
def test_resumed_from_in_agent_line(self) -> None:
|
||||
self.mgr.start_agent(
|
||||
prompt="t", label="res",
|
||||
resumed_from_session_id="old-sess-id",
|
||||
)
|
||||
lines = self.mgr.summary_lines()
|
||||
detail = [l for l in lines if "res" in l]
|
||||
self.assertTrue(any("resumed_from=old-sess-id" in l for l in detail))
|
||||
|
||||
def test_agent_without_label_uses_id(self) -> None:
|
||||
aid = self.mgr.start_agent(prompt="no label")
|
||||
lines = self.mgr.summary_lines()
|
||||
detail = [l for l in lines if aid in l]
|
||||
self.assertEqual(len(detail), 1)
|
||||
|
||||
def test_truncation_at_8_agents(self) -> None:
|
||||
for i in range(10):
|
||||
self.mgr.start_agent(prompt=f"task-{i}")
|
||||
lines = self.mgr.summary_lines()
|
||||
self.assertIn("- Managed agents: 10", lines)
|
||||
plus_line = [l for l in lines if "plus" in l and "managed agents" in l]
|
||||
self.assertEqual(len(plus_line), 1)
|
||||
self.assertIn("2 more managed agents", plus_line[0])
|
||||
|
||||
def test_truncation_at_6_groups(self) -> None:
|
||||
for i in range(8):
|
||||
self.mgr.start_group(label=f"grp-{i}")
|
||||
lines = self.mgr.summary_lines()
|
||||
plus_line = [l for l in lines if "plus" in l and "agent groups" in l]
|
||||
self.assertEqual(len(plus_line), 1)
|
||||
self.assertIn("2 more agent groups", plus_line[0])
|
||||
|
||||
def test_group_summary_line_format(self) -> None:
|
||||
gid = self.mgr.start_group(label="my-batch", strategy="parallel")
|
||||
a1 = self.mgr.start_agent(prompt="t1", group_id=gid, child_index=0)
|
||||
self.mgr.finish_agent(
|
||||
a1, session_id="s", session_path="/p",
|
||||
turns=1, tool_calls=2, stop_reason="end_turn",
|
||||
)
|
||||
self.mgr.finish_group(
|
||||
gid, status="completed",
|
||||
completed_children=1, failed_children=0,
|
||||
batch_count=1, max_batch_size=1,
|
||||
)
|
||||
lines = self.mgr.summary_lines()
|
||||
grp_line = [l for l in lines if "my-batch" in l and "group_status" in l]
|
||||
self.assertEqual(len(grp_line), 1)
|
||||
self.assertIn("group_status=completed", grp_line[0])
|
||||
self.assertIn("children=1", grp_line[0])
|
||||
self.assertIn("completed=1", grp_line[0])
|
||||
self.assertIn("failed=0", grp_line[0])
|
||||
self.assertIn("strategy=parallel", grp_line[0])
|
||||
self.assertIn("stop_reasons=end_turn:1", grp_line[0])
|
||||
|
||||
def test_child_and_resumed_counts(self) -> None:
|
||||
p = self.mgr.start_agent(prompt="parent")
|
||||
self.mgr.start_agent(prompt="c1", parent_agent_id=p)
|
||||
self.mgr.start_agent(
|
||||
prompt="c2", parent_agent_id=p,
|
||||
resumed_from_session_id="old",
|
||||
)
|
||||
lines = self.mgr.summary_lines()
|
||||
self.assertIn("- Child agents: 2", lines)
|
||||
self.assertIn("- Resumed agents: 1", lines)
|
||||
|
||||
|
||||
class TestMultipleAgentsAndGroupsInteraction(unittest.TestCase):
|
||||
"""End-to-end scenario with multiple groups and cross-references."""
|
||||
|
||||
def test_full_lifecycle(self) -> None:
|
||||
mgr = AgentManager()
|
||||
|
||||
# Parent agent spawns two groups
|
||||
parent = mgr.start_agent(prompt="orchestrate", label="orchestrator")
|
||||
g1 = mgr.start_group(label="build", parent_agent_id=parent, strategy="serial")
|
||||
g2 = mgr.start_group(label="test", parent_agent_id=parent, strategy="parallel")
|
||||
|
||||
# Group 1 children
|
||||
b1 = mgr.start_agent(prompt="build-fe", group_id=g1, child_index=0, parent_agent_id=parent)
|
||||
b2 = mgr.start_agent(prompt="build-be", group_id=g1, child_index=1, parent_agent_id=parent)
|
||||
|
||||
# Group 2 children, one resumed
|
||||
t1 = mgr.start_agent(
|
||||
prompt="test-unit", group_id=g2, child_index=0,
|
||||
parent_agent_id=parent, resumed_from_session_id="old-session",
|
||||
)
|
||||
t2 = mgr.start_agent(prompt="test-e2e", group_id=g2, child_index=1, parent_agent_id=parent)
|
||||
|
||||
# Finish agents
|
||||
for aid, turns, tc, sr in [
|
||||
(b1, 3, 8, "end_turn"),
|
||||
(b2, 4, 10, "end_turn"),
|
||||
(t1, 2, 5, "end_turn"),
|
||||
(t2, 6, 15, "max_turns"),
|
||||
]:
|
||||
mgr.finish_agent(aid, session_id=f"s-{aid}", session_path=f"/p/{aid}", turns=turns, tool_calls=tc, stop_reason=sr)
|
||||
|
||||
mgr.finish_group(g1, status="completed", completed_children=2, failed_children=0, batch_count=2, max_batch_size=1)
|
||||
mgr.finish_group(g2, status="completed", completed_children=1, failed_children=1, batch_count=1, max_batch_size=2, dependency_skips=1)
|
||||
|
||||
# Verify children_of
|
||||
children = mgr.children_of(parent)
|
||||
self.assertEqual(len(children), 4)
|
||||
|
||||
# Verify group_children ordering
|
||||
g1_children = mgr.group_children(g1)
|
||||
self.assertEqual(g1_children[0].agent_id, b1)
|
||||
self.assertEqual(g1_children[1].agent_id, b2)
|
||||
|
||||
g2_children = mgr.group_children(g2)
|
||||
self.assertEqual(g2_children[0].agent_id, t1)
|
||||
self.assertEqual(g2_children[1].agent_id, t2)
|
||||
|
||||
# Verify completed records (parent is still running)
|
||||
completed = mgr.completed_records()
|
||||
self.assertEqual(len(completed), 4)
|
||||
|
||||
# Verify group summaries
|
||||
s1 = mgr.group_summary(g1)
|
||||
assert s1 is not None
|
||||
self.assertEqual(s1["child_count"], 2)
|
||||
self.assertEqual(s1["resumed_children"], 0)
|
||||
self.assertEqual(s1["dependency_skips"], 0)
|
||||
|
||||
s2 = mgr.group_summary(g2)
|
||||
assert s2 is not None
|
||||
self.assertEqual(s2["child_count"], 2)
|
||||
self.assertEqual(s2["resumed_children"], 1)
|
||||
self.assertEqual(s2["dependency_skips"], 1)
|
||||
self.assertEqual(s2["stop_reason_counts"], {"end_turn": 1, "max_turns": 1})
|
||||
|
||||
# Verify summary_lines produces output
|
||||
lines = mgr.summary_lines()
|
||||
self.assertIn("- Managed agents: 5", lines)
|
||||
self.assertIn("- Completed agents: 4", lines)
|
||||
self.assertIn("- Child agents: 4", lines)
|
||||
self.assertIn("- Resumed agents: 1", lines)
|
||||
self.assertIn("- Agent groups: 2", lines)
|
||||
self.assertIn("- Completed groups: 2", lines)
|
||||
|
||||
|
||||
class TestFrozenDataclasses(unittest.TestCase):
|
||||
def test_record_is_frozen(self) -> None:
|
||||
rec = ManagedAgentRecord(agent_id="a", prompt="p")
|
||||
with self.assertRaises(AttributeError):
|
||||
rec.status = "completed" # type: ignore[misc]
|
||||
|
||||
def test_group_is_frozen(self) -> None:
|
||||
grp = ManagedAgentGroup(group_id="g")
|
||||
with self.assertRaises(AttributeError):
|
||||
grp.status = "completed" # type: ignore[misc]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,475 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from src.agent_plugin_cache import (
|
||||
MAX_PLUGIN_LINES,
|
||||
MAX_PLUGIN_PREVIEW_CHARS,
|
||||
PluginCacheEntry,
|
||||
_coerce_entry,
|
||||
_extract_entries,
|
||||
discover_plugin_cache,
|
||||
load_plugin_cache_summary,
|
||||
)
|
||||
|
||||
|
||||
class TestDiscoverPluginCacheNone(unittest.TestCase):
|
||||
"""discover_plugin_cache returns None when no cache files exist."""
|
||||
|
||||
def test_returns_none_for_empty_directory(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_returns_none_when_port_sessions_dir_is_empty(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
(Path(tmp) / ".port_sessions").mkdir()
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestDiscoverPluginCacheListFormat(unittest.TestCase):
|
||||
"""discover_plugin_cache finds cache in .port_sessions/plugin_cache.json (list format)."""
|
||||
|
||||
def test_list_of_strings(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "plugin_cache.json"
|
||||
cache_file.write_text(json.dumps(["plugin-a", "plugin-b"]))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("plugin-a", result)
|
||||
self.assertIn("plugin-b", result)
|
||||
self.assertIn("Plugin entries discovered: 2", result)
|
||||
|
||||
def test_list_of_dicts(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "plugin_cache.json"
|
||||
cache_file.write_text(
|
||||
json.dumps([{"name": "alpha", "version": "1.0"}, {"name": "beta"}])
|
||||
)
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("alpha", result)
|
||||
self.assertIn("version=1.0", result)
|
||||
self.assertIn("beta", result)
|
||||
|
||||
|
||||
class TestDiscoverPluginCacheDictPluginsKey(unittest.TestCase):
|
||||
"""discover_plugin_cache finds cache in .port_sessions/plugins.json (dict with 'plugins' key)."""
|
||||
|
||||
def test_dict_with_plugins_list(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "plugins.json"
|
||||
payload = {"plugins": [{"name": "foo"}, {"name": "bar"}]}
|
||||
cache_file.write_text(json.dumps(payload))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("foo", result)
|
||||
self.assertIn("bar", result)
|
||||
self.assertIn("Plugin entries discovered: 2", result)
|
||||
|
||||
|
||||
class TestDiscoverPluginCacheDictEntriesKey(unittest.TestCase):
|
||||
"""discover_plugin_cache handles dict with 'entries' key format."""
|
||||
|
||||
def test_dict_with_entries_list(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "plugin_cache.json"
|
||||
payload = {"entries": [{"name": "entry-a"}, {"name": "entry-b"}]}
|
||||
cache_file.write_text(json.dumps(payload))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("entry-a", result)
|
||||
self.assertIn("entry-b", result)
|
||||
|
||||
|
||||
class TestDiscoverPluginCacheDictKeyAsName(unittest.TestCase):
|
||||
"""discover_plugin_cache handles dict where values are dicts (key=name format)."""
|
||||
|
||||
def test_dict_values_are_dicts(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "plugin_cache.json"
|
||||
payload = {
|
||||
"my-plugin": {"version": "2.0", "source": "/path/to/it"},
|
||||
"other-plugin": {"version": "3.1"},
|
||||
}
|
||||
cache_file.write_text(json.dumps(payload))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("my-plugin", result)
|
||||
self.assertIn("version=2.0", result)
|
||||
self.assertIn("source=/path/to/it", result)
|
||||
self.assertIn("other-plugin", result)
|
||||
|
||||
|
||||
class TestCoerceEntry(unittest.TestCase):
|
||||
"""_coerce_entry handles various input types."""
|
||||
|
||||
def test_string_entry(self) -> None:
|
||||
entry = _coerce_entry("simple-plugin")
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "simple-plugin")
|
||||
self.assertTrue(entry.enabled)
|
||||
|
||||
def test_string_entry_strips_whitespace(self) -> None:
|
||||
entry = _coerce_entry(" padded-name ")
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "padded-name")
|
||||
|
||||
def test_empty_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_coerce_entry(""))
|
||||
self.assertIsNone(_coerce_entry(" "))
|
||||
|
||||
def test_dict_with_name_key(self) -> None:
|
||||
entry = _coerce_entry({"name": "named-plugin"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "named-plugin")
|
||||
|
||||
def test_dict_with_plugin_key(self) -> None:
|
||||
entry = _coerce_entry({"plugin": "plugin-key"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "plugin-key")
|
||||
|
||||
def test_dict_with_id_key(self) -> None:
|
||||
entry = _coerce_entry({"id": "id-key"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "id-key")
|
||||
|
||||
def test_name_takes_precedence_over_plugin_and_id(self) -> None:
|
||||
entry = _coerce_entry({"name": "winner", "plugin": "loser", "id": "also-loser"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.name, "winner")
|
||||
|
||||
def test_dict_with_version_and_source(self) -> None:
|
||||
entry = _coerce_entry(
|
||||
{"name": "full", "version": "1.2.3", "source": "/src"}
|
||||
)
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.version, "1.2.3")
|
||||
self.assertEqual(entry.source, "/src")
|
||||
|
||||
def test_source_fallback_to_path(self) -> None:
|
||||
entry = _coerce_entry({"name": "p", "path": "/a/b"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.source, "/a/b")
|
||||
|
||||
def test_source_fallback_to_module(self) -> None:
|
||||
entry = _coerce_entry({"name": "p", "module": "my.mod"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertEqual(entry.source, "my.mod")
|
||||
|
||||
def test_disabled_plugin(self) -> None:
|
||||
entry = _coerce_entry({"name": "off", "enabled": False})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertFalse(entry.enabled)
|
||||
|
||||
def test_enabled_none_defaults_to_true(self) -> None:
|
||||
entry = _coerce_entry({"name": "on"})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertTrue(entry.enabled)
|
||||
|
||||
def test_enabled_truthy_value(self) -> None:
|
||||
entry = _coerce_entry({"name": "on", "enabled": 1})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertTrue(entry.enabled)
|
||||
|
||||
def test_empty_dict_returns_none(self) -> None:
|
||||
self.assertIsNone(_coerce_entry({}))
|
||||
|
||||
def test_non_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_coerce_entry(42))
|
||||
self.assertIsNone(_coerce_entry(None))
|
||||
self.assertIsNone(_coerce_entry(True))
|
||||
self.assertIsNone(_coerce_entry([]))
|
||||
|
||||
def test_dict_with_non_string_name_returns_none(self) -> None:
|
||||
self.assertIsNone(_coerce_entry({"name": 123}))
|
||||
self.assertIsNone(_coerce_entry({"name": ""}))
|
||||
|
||||
def test_empty_version_is_none(self) -> None:
|
||||
entry = _coerce_entry({"name": "p", "version": ""})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertIsNone(entry.version)
|
||||
|
||||
def test_non_string_version_is_none(self) -> None:
|
||||
entry = _coerce_entry({"name": "p", "version": 5})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertIsNone(entry.version)
|
||||
|
||||
def test_empty_source_is_none(self) -> None:
|
||||
entry = _coerce_entry({"name": "p", "source": ""})
|
||||
self.assertIsNotNone(entry)
|
||||
self.assertIsNone(entry.source)
|
||||
|
||||
|
||||
class TestExtractEntries(unittest.TestCase):
|
||||
"""_extract_entries handles all payload shapes."""
|
||||
|
||||
def test_list_payload(self) -> None:
|
||||
entries = _extract_entries(["a", "b"])
|
||||
self.assertEqual(len(entries), 2)
|
||||
|
||||
def test_dict_plugins_key(self) -> None:
|
||||
entries = _extract_entries({"plugins": [{"name": "x"}]})
|
||||
self.assertEqual(len(entries), 1)
|
||||
self.assertEqual(entries[0].name, "x")
|
||||
|
||||
def test_dict_entries_key(self) -> None:
|
||||
entries = _extract_entries({"entries": [{"name": "y"}]})
|
||||
self.assertEqual(len(entries), 1)
|
||||
self.assertEqual(entries[0].name, "y")
|
||||
|
||||
def test_dict_key_as_name(self) -> None:
|
||||
entries = _extract_entries({"k1": {"version": "1"}, "k2": {"version": "2"}})
|
||||
names = {e.name for e in entries}
|
||||
self.assertEqual(names, {"k1", "k2"})
|
||||
|
||||
def test_plugins_key_takes_precedence_over_key_as_name(self) -> None:
|
||||
payload = {"plugins": [{"name": "from-plugins"}], "other": {"version": "1"}}
|
||||
entries = _extract_entries(payload)
|
||||
self.assertEqual(len(entries), 1)
|
||||
self.assertEqual(entries[0].name, "from-plugins")
|
||||
|
||||
def test_non_dict_values_ignored_in_key_as_name(self) -> None:
|
||||
entries = _extract_entries({"good": {"version": "1"}, "bad": "string-val"})
|
||||
self.assertEqual(len(entries), 1)
|
||||
self.assertEqual(entries[0].name, "good")
|
||||
|
||||
def test_empty_list_returns_empty(self) -> None:
|
||||
self.assertEqual(_extract_entries([]), [])
|
||||
|
||||
def test_invalid_payload_type(self) -> None:
|
||||
self.assertEqual(_extract_entries("not-valid"), [])
|
||||
self.assertEqual(_extract_entries(42), [])
|
||||
|
||||
|
||||
class TestLoadPluginCacheSummary(unittest.TestCase):
|
||||
"""load_plugin_cache_summary returns rendered summary string."""
|
||||
|
||||
def test_returns_none_when_no_cache(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = load_plugin_cache_summary(Path(tmp))
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_returns_summary_string(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text(
|
||||
json.dumps([{"name": "my-plugin", "version": "1.0"}])
|
||||
)
|
||||
|
||||
result = load_plugin_cache_summary(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("my-plugin", result)
|
||||
self.assertIn("Plugin cache loaded from:", result)
|
||||
|
||||
|
||||
class TestRenderedSummaryCounts(unittest.TestCase):
|
||||
"""Rendered summary shows correct enabled/disabled counts."""
|
||||
|
||||
def _make_cache(self, tmp: str, entries: list) -> str | None:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir(exist_ok=True)
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps(entries))
|
||||
return discover_plugin_cache(Path(tmp))
|
||||
|
||||
def test_all_enabled(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = self._make_cache(tmp, [{"name": "a"}, {"name": "b"}, {"name": "c"}])
|
||||
self.assertIn("Enabled plugins: 3", result)
|
||||
self.assertNotIn("Disabled plugins:", result)
|
||||
|
||||
def test_some_disabled(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = self._make_cache(
|
||||
tmp,
|
||||
[
|
||||
{"name": "a"},
|
||||
{"name": "b", "enabled": False},
|
||||
{"name": "c", "enabled": False},
|
||||
],
|
||||
)
|
||||
self.assertIn("Enabled plugins: 1", result)
|
||||
self.assertIn("Disabled plugins: 2", result)
|
||||
|
||||
def test_disabled_shown_in_line(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = self._make_cache(
|
||||
tmp, [{"name": "off-plugin", "enabled": False}]
|
||||
)
|
||||
self.assertIn("disabled", result)
|
||||
self.assertIn("off-plugin", result)
|
||||
|
||||
|
||||
class TestPreviewTruncation(unittest.TestCase):
|
||||
"""Preview truncation works (MAX_PLUGIN_PREVIEW_CHARS=4000)."""
|
||||
|
||||
def test_long_output_is_truncated(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
# Create entries with very long names so the rendered output exceeds the limit
|
||||
long_name = "x" * 500
|
||||
entries = [{"name": f"{long_name}-{i}"} for i in range(20)]
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps(entries))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertLessEqual(len(result), MAX_PLUGIN_PREVIEW_CHARS)
|
||||
self.assertTrue(result.endswith("..."))
|
||||
|
||||
def test_short_output_not_truncated(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
entries = [{"name": "small"}]
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps(entries))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertFalse(result.endswith("..."))
|
||||
|
||||
|
||||
class TestMaxPluginLinesTruncation(unittest.TestCase):
|
||||
"""More than MAX_PLUGIN_LINES (12) shows truncation message."""
|
||||
|
||||
def test_more_than_max_lines_shows_truncation(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
count = MAX_PLUGIN_LINES + 5
|
||||
entries = [{"name": f"plugin-{i}"} for i in range(count)]
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps(entries))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn(f"... plus 5 more plugin entries", result)
|
||||
self.assertIn(f"Plugin entries discovered: {count}", result)
|
||||
|
||||
def test_exactly_max_lines_no_truncation(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
entries = [{"name": f"plugin-{i}"} for i in range(MAX_PLUGIN_LINES)]
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps(entries))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertNotIn("more plugin entries", result)
|
||||
|
||||
|
||||
class TestMalformedJsonSkipped(unittest.TestCase):
|
||||
"""Malformed JSON files are gracefully skipped."""
|
||||
|
||||
def test_invalid_json_skipped(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text("{not valid json!!!")
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_malformed_first_valid_second(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
# First candidate: malformed
|
||||
(cache_dir / "plugin_cache.json").write_text("not json")
|
||||
# Second candidate: valid
|
||||
(cache_dir / "plugins.json").write_text(
|
||||
json.dumps({"plugins": [{"name": "fallback"}]})
|
||||
)
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("fallback", result)
|
||||
|
||||
def test_valid_json_but_empty_entries(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
cache_dir = Path(tmp) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text(json.dumps([]))
|
||||
|
||||
result = discover_plugin_cache(Path(tmp))
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestAdditionalWorkingDirectories(unittest.TestCase):
|
||||
"""additional_working_directories are searched."""
|
||||
|
||||
def test_finds_cache_in_additional_dir(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as main, tempfile.TemporaryDirectory() as extra:
|
||||
cache_dir = Path(extra) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text(
|
||||
json.dumps([{"name": "extra-plugin"}])
|
||||
)
|
||||
|
||||
result = discover_plugin_cache(Path(main), (extra,))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("extra-plugin", result)
|
||||
|
||||
def test_main_dir_preferred_over_additional(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as main, tempfile.TemporaryDirectory() as extra:
|
||||
for base, name in [(main, "main-plugin"), (extra, "extra-plugin")]:
|
||||
cache_dir = Path(base) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text(
|
||||
json.dumps([{"name": name}])
|
||||
)
|
||||
|
||||
result = discover_plugin_cache(Path(main), (extra,))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("main-plugin", result)
|
||||
|
||||
def test_load_plugin_cache_summary_with_additional_dirs(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as main, tempfile.TemporaryDirectory() as extra:
|
||||
cache_dir = Path(extra) / ".port_sessions"
|
||||
cache_dir.mkdir()
|
||||
(cache_dir / "plugin_cache.json").write_text(
|
||||
json.dumps([{"name": "via-summary"}])
|
||||
)
|
||||
|
||||
result = load_plugin_cache_summary(Path(main), (extra,))
|
||||
self.assertIsNotNone(result)
|
||||
self.assertIn("via-summary", result)
|
||||
|
||||
|
||||
class TestPluginCacheEntry(unittest.TestCase):
|
||||
"""PluginCacheEntry dataclass behavior."""
|
||||
|
||||
def test_defaults(self) -> None:
|
||||
entry = PluginCacheEntry(name="test")
|
||||
self.assertEqual(entry.name, "test")
|
||||
self.assertTrue(entry.enabled)
|
||||
self.assertIsNone(entry.version)
|
||||
self.assertIsNone(entry.source)
|
||||
|
||||
def test_frozen(self) -> None:
|
||||
entry = PluginCacheEntry(name="test")
|
||||
with self.assertRaises(AttributeError):
|
||||
entry.name = "other" # type: ignore[misc]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Security tests for agent_tools.py: path traversal, destructive commands, and env var filtering."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from src.agent_tools import (
|
||||
ToolExecutionError,
|
||||
ToolPermissionError,
|
||||
_ensure_shell_allowed,
|
||||
_is_sensitive_env_var,
|
||||
_resolve_path,
|
||||
build_tool_context,
|
||||
default_tool_registry,
|
||||
)
|
||||
from src.agent_types import AgentPermissions, AgentRuntimeConfig
|
||||
|
||||
|
||||
def _make_context(
|
||||
tmp_dir: str,
|
||||
*,
|
||||
allow_shell: bool = False,
|
||||
allow_destructive: bool = False,
|
||||
) -> "ToolExecutionContext": # noqa: F821
|
||||
config = AgentRuntimeConfig(
|
||||
cwd=Path(tmp_dir),
|
||||
permissions=AgentPermissions(
|
||||
allow_shell_commands=allow_shell,
|
||||
allow_destructive_shell_commands=allow_destructive,
|
||||
),
|
||||
)
|
||||
return build_tool_context(config, tool_registry=default_tool_registry())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _resolve_path – path traversal prevention
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestResolvePath(unittest.TestCase):
|
||||
def test_relative_path_within_workspace(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
(Path(tmp) / "hello.txt").write_text("hi")
|
||||
ctx = _make_context(tmp)
|
||||
result = _resolve_path("hello.txt", ctx)
|
||||
self.assertEqual(result, (Path(tmp) / "hello.txt").resolve())
|
||||
|
||||
def test_absolute_path_within_workspace(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
target = Path(tmp) / "sub" / "file.txt"
|
||||
target.parent.mkdir()
|
||||
target.write_text("data")
|
||||
ctx = _make_context(tmp)
|
||||
result = _resolve_path(str(target), ctx)
|
||||
self.assertEqual(result, target.resolve())
|
||||
|
||||
def test_traversal_with_dotdot_blocked(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
ctx = _make_context(tmp)
|
||||
with self.assertRaises(ToolExecutionError):
|
||||
_resolve_path("../outside", ctx)
|
||||
|
||||
def test_traversal_etc_passwd_blocked(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
ctx = _make_context(tmp)
|
||||
with self.assertRaises(ToolExecutionError):
|
||||
_resolve_path("../../etc/passwd", ctx)
|
||||
|
||||
def test_allow_missing_true_permits_nonexistent(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
ctx = _make_context(tmp)
|
||||
result = _resolve_path("does_not_exist.txt", ctx, allow_missing=True)
|
||||
self.assertEqual(result, (Path(tmp) / "does_not_exist.txt").resolve())
|
||||
|
||||
def test_allow_missing_false_raises_for_nonexistent(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
ctx = _make_context(tmp)
|
||||
with self.assertRaises(OSError):
|
||||
_resolve_path("does_not_exist.txt", ctx, allow_missing=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _ensure_shell_allowed – destructive command blocking
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestEnsureShellAllowed(unittest.TestCase):
|
||||
def _ctx(self, *, allow_shell: bool = True, allow_destructive: bool = False) -> "ToolExecutionContext": # noqa: F821
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
return _make_context(
|
||||
self._tmp.name,
|
||||
allow_shell=allow_shell,
|
||||
allow_destructive=allow_destructive,
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
if hasattr(self, "_tmp"):
|
||||
self._tmp.cleanup()
|
||||
|
||||
# -- safe commands pass --------------------------------------------------
|
||||
def test_safe_commands_allowed(self):
|
||||
ctx = self._ctx()
|
||||
for cmd in ("ls -la", "cat file.txt", "echo hello", "grep foo bar.txt"):
|
||||
_ensure_shell_allowed(cmd, ctx) # should not raise
|
||||
|
||||
# -- destructive commands blocked -----------------------------------------
|
||||
def test_rm_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("rm -rf /", ctx)
|
||||
|
||||
def test_mv_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("mv a b", ctx)
|
||||
|
||||
def test_dd_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("dd if=/dev/zero of=/dev/sda", ctx)
|
||||
|
||||
def test_shutdown_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("shutdown -h now", ctx)
|
||||
|
||||
def test_reboot_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("reboot ", ctx)
|
||||
|
||||
def test_mkfs_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("mkfs.ext4 /dev/sda1", ctx)
|
||||
|
||||
def test_chmod_recursive_777_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("chmod -R 777 /", ctx)
|
||||
|
||||
def test_chown_recursive_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("chown -R root:root /", ctx)
|
||||
|
||||
def test_git_reset_hard_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("git reset --hard", ctx)
|
||||
|
||||
def test_git_clean_fd_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("git clean -fd", ctx)
|
||||
|
||||
def test_truncation_operator_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed(": > important.log", ctx)
|
||||
|
||||
# -- chained commands with destructive sub-commands -----------------------
|
||||
def test_chained_and_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("echo hi && rm -rf /", ctx)
|
||||
|
||||
def test_chained_or_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("false || rm file", ctx)
|
||||
|
||||
def test_chained_semicolon_blocked(self):
|
||||
ctx = self._ctx()
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("echo ok; mv a b", ctx)
|
||||
|
||||
# -- shell commands entirely disabled ------------------------------------
|
||||
def test_shell_disabled_raises(self):
|
||||
ctx = self._ctx(allow_shell=False)
|
||||
with self.assertRaises(ToolPermissionError):
|
||||
_ensure_shell_allowed("ls", ctx)
|
||||
|
||||
# -- allow_destructive bypasses blocking ---------------------------------
|
||||
def test_destructive_allowed_bypasses(self):
|
||||
ctx = self._ctx(allow_destructive=True)
|
||||
# All destructive commands should pass without raising
|
||||
for cmd in (
|
||||
"rm -rf /",
|
||||
"mv a b",
|
||||
"dd if=/dev/zero of=/dev/sda",
|
||||
"shutdown -h now",
|
||||
"mkfs.ext4 /dev/sda1",
|
||||
"chmod -R 777 /",
|
||||
"chown -R root:root /",
|
||||
"git reset --hard",
|
||||
"git clean -fd",
|
||||
": > file",
|
||||
):
|
||||
_ensure_shell_allowed(cmd, ctx) # should not raise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _is_sensitive_env_var – secret-name detection
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestIsSensitiveEnvVar(unittest.TestCase):
|
||||
def test_common_sensitive_vars_detected(self):
|
||||
for name in (
|
||||
"MY_SECRET",
|
||||
"GITHUB_TOKEN",
|
||||
"DB_PASSWORD",
|
||||
"SSH_PRIVATE_KEY",
|
||||
"MY_API_KEY",
|
||||
"CREDENTIAL_STORE",
|
||||
"AUTH_HEADER",
|
||||
):
|
||||
self.assertTrue(
|
||||
_is_sensitive_env_var(name),
|
||||
f"{name} should be detected as sensitive",
|
||||
)
|
||||
|
||||
def test_non_sensitive_vars_allowed(self):
|
||||
for name in ("HOME", "PATH", "LANG", "TERM", "USER", "SHELL"):
|
||||
self.assertFalse(
|
||||
_is_sensitive_env_var(name),
|
||||
f"{name} should not be detected as sensitive",
|
||||
)
|
||||
|
||||
def test_case_insensitive(self):
|
||||
self.assertTrue(_is_sensitive_env_var("my_secret"))
|
||||
self.assertTrue(_is_sensitive_env_var("Github_Token"))
|
||||
self.assertTrue(_is_sensitive_env_var("db_password"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,198 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from src.command_graph import CommandGraph, build_command_graph
|
||||
from src.models import PortingModule
|
||||
|
||||
|
||||
def _module(name: str, source_hint: str) -> PortingModule:
|
||||
return PortingModule(name=name, responsibility="stub", source_hint=source_hint)
|
||||
|
||||
|
||||
class CommandGraphTests(unittest.TestCase):
|
||||
# -- construction & immutability ------------------------------------------
|
||||
|
||||
def test_empty_graph(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
self.assertEqual(graph.builtins, ())
|
||||
self.assertEqual(graph.plugin_like, ())
|
||||
self.assertEqual(graph.skill_like, ())
|
||||
|
||||
def test_fields_are_tuples(self) -> None:
|
||||
b = (_module("b1", "core/b1.ts"),)
|
||||
p = (_module("p1", "plugin/p1.ts"),)
|
||||
s = (_module("s1", "skills/s1.ts"),)
|
||||
graph = CommandGraph(builtins=b, plugin_like=p, skill_like=s)
|
||||
self.assertIsInstance(graph.builtins, tuple)
|
||||
self.assertIsInstance(graph.plugin_like, tuple)
|
||||
self.assertIsInstance(graph.skill_like, tuple)
|
||||
|
||||
def test_frozen_dataclass_rejects_mutation(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
with self.assertRaises(AttributeError):
|
||||
graph.builtins = () # type: ignore[misc]
|
||||
|
||||
# -- flattened -------------------------------------------------------------
|
||||
|
||||
def test_flattened_combines_all_categories(self) -> None:
|
||||
b = (_module("b1", "core/b1.ts"),)
|
||||
p = (_module("p1", "plugin/p1.ts"),)
|
||||
s = (_module("s1", "skills/s1.ts"),)
|
||||
graph = CommandGraph(builtins=b, plugin_like=p, skill_like=s)
|
||||
self.assertEqual(graph.flattened(), b + p + s)
|
||||
|
||||
def test_flattened_preserves_order(self) -> None:
|
||||
b1 = _module("b1", "core/b1.ts")
|
||||
b2 = _module("b2", "core/b2.ts")
|
||||
p1 = _module("p1", "plugin/p1.ts")
|
||||
s1 = _module("s1", "skills/s1.ts")
|
||||
graph = CommandGraph(builtins=(b1, b2), plugin_like=(p1,), skill_like=(s1,))
|
||||
self.assertEqual(graph.flattened(), (b1, b2, p1, s1))
|
||||
|
||||
def test_flattened_of_empty_graph_returns_empty_tuple(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
self.assertEqual(graph.flattened(), ())
|
||||
|
||||
def test_flattened_length_is_sum_of_categories(self) -> None:
|
||||
b = (_module("b1", "core/b1.ts"), _module("b2", "core/b2.ts"))
|
||||
p = (_module("p1", "plugin/p1.ts"),)
|
||||
s = ()
|
||||
graph = CommandGraph(builtins=b, plugin_like=p, skill_like=s)
|
||||
self.assertEqual(len(graph.flattened()), 3)
|
||||
|
||||
# -- as_markdown -----------------------------------------------------------
|
||||
|
||||
def test_as_markdown_includes_header(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
md = graph.as_markdown()
|
||||
self.assertIn("# Command Graph", md)
|
||||
|
||||
def test_as_markdown_includes_counts(self) -> None:
|
||||
b = (_module("b1", "core/b1.ts"), _module("b2", "core/b2.ts"))
|
||||
p = (_module("p1", "plugin/p1.ts"),)
|
||||
s = (_module("s1", "skills/s1.ts"), _module("s2", "skills/s2.ts"), _module("s3", "skills/s3.ts"))
|
||||
graph = CommandGraph(builtins=b, plugin_like=p, skill_like=s)
|
||||
md = graph.as_markdown()
|
||||
self.assertIn("Builtins: 2", md)
|
||||
self.assertIn("Plugin-like commands: 1", md)
|
||||
self.assertIn("Skill-like commands: 3", md)
|
||||
|
||||
def test_as_markdown_empty_counts(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
md = graph.as_markdown()
|
||||
self.assertIn("Builtins: 0", md)
|
||||
self.assertIn("Plugin-like commands: 0", md)
|
||||
self.assertIn("Skill-like commands: 0", md)
|
||||
|
||||
def test_as_markdown_returns_string(self) -> None:
|
||||
graph = CommandGraph(builtins=(), plugin_like=(), skill_like=())
|
||||
self.assertIsInstance(graph.as_markdown(), str)
|
||||
|
||||
|
||||
class BuildCommandGraphTests(unittest.TestCase):
|
||||
# -- return type -----------------------------------------------------------
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_returns_command_graph(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
mock_get.return_value = ()
|
||||
result = build_command_graph()
|
||||
self.assertIsInstance(result, CommandGraph)
|
||||
|
||||
# -- categorization --------------------------------------------------------
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_plugin_source_goes_to_plugin_like(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
p = _module("p1", "plugin/p1.ts")
|
||||
mock_get.return_value = (p,)
|
||||
graph = build_command_graph()
|
||||
self.assertIn(p, graph.plugin_like)
|
||||
self.assertNotIn(p, graph.builtins)
|
||||
self.assertNotIn(p, graph.skill_like)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_skills_source_goes_to_skill_like(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
s = _module("s1", "skills/s1.ts")
|
||||
mock_get.return_value = (s,)
|
||||
graph = build_command_graph()
|
||||
self.assertIn(s, graph.skill_like)
|
||||
self.assertNotIn(s, graph.builtins)
|
||||
self.assertNotIn(s, graph.plugin_like)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_plain_source_goes_to_builtins(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
b = _module("b1", "core/b1.ts")
|
||||
mock_get.return_value = (b,)
|
||||
graph = build_command_graph()
|
||||
self.assertIn(b, graph.builtins)
|
||||
self.assertNotIn(b, graph.plugin_like)
|
||||
self.assertNotIn(b, graph.skill_like)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_case_insensitive_plugin_match(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
p = _module("p1", "Plugin/p1.ts")
|
||||
mock_get.return_value = (p,)
|
||||
graph = build_command_graph()
|
||||
self.assertIn(p, graph.plugin_like)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_case_insensitive_skills_match(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
s = _module("s1", "Skills/s1.ts")
|
||||
mock_get.return_value = (s,)
|
||||
graph = build_command_graph()
|
||||
self.assertIn(s, graph.skill_like)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_mixed_commands_are_sorted_correctly(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
b = _module("b1", "core/b1.ts")
|
||||
p = _module("p1", "plugin/p1.ts")
|
||||
s = _module("s1", "skills/s1.ts")
|
||||
mock_get.return_value = (b, p, s)
|
||||
graph = build_command_graph()
|
||||
self.assertEqual(graph.builtins, (b,))
|
||||
self.assertEqual(graph.plugin_like, (p,))
|
||||
self.assertEqual(graph.skill_like, (s,))
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_empty_commands_yields_empty_graph(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
mock_get.return_value = ()
|
||||
graph = build_command_graph()
|
||||
self.assertEqual(graph.builtins, ())
|
||||
self.assertEqual(graph.plugin_like, ())
|
||||
self.assertEqual(graph.skill_like, ())
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_all_builtins(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
b1 = _module("b1", "core/b1.ts")
|
||||
b2 = _module("b2", "commands/b2.ts")
|
||||
mock_get.return_value = (b1, b2)
|
||||
graph = build_command_graph()
|
||||
self.assertEqual(len(graph.builtins), 2)
|
||||
self.assertEqual(graph.plugin_like, ())
|
||||
self.assertEqual(graph.skill_like, ())
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_builtins_plugin_skill_are_tuples_of_porting_module(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
b = _module("b1", "core/b1.ts")
|
||||
p = _module("p1", "plugin/p1.ts")
|
||||
s = _module("s1", "skills/s1.ts")
|
||||
mock_get.return_value = (b, p, s)
|
||||
graph = build_command_graph()
|
||||
for category in (graph.builtins, graph.plugin_like, graph.skill_like):
|
||||
self.assertIsInstance(category, tuple)
|
||||
for item in category:
|
||||
self.assertIsInstance(item, PortingModule)
|
||||
|
||||
@patch("src.command_graph.get_commands")
|
||||
def test_flattened_matches_original_commands(self, mock_get: unittest.mock.MagicMock) -> None:
|
||||
modules = (
|
||||
_module("b1", "core/b1.ts"),
|
||||
_module("p1", "plugin/p1.ts"),
|
||||
_module("s1", "skills/s1.ts"),
|
||||
)
|
||||
mock_get.return_value = modules
|
||||
graph = build_command_graph()
|
||||
self.assertEqual(set(graph.flattened()), set(modules))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,52 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.cost_tracker import CostTracker
|
||||
|
||||
|
||||
class CostTrackerTests(unittest.TestCase):
|
||||
def test_fresh_tracker_starts_at_zero_with_empty_events(self) -> None:
|
||||
tracker = CostTracker()
|
||||
self.assertEqual(tracker.total_units, 0)
|
||||
self.assertEqual(tracker.events, [])
|
||||
|
||||
def test_record_single_event(self) -> None:
|
||||
tracker = CostTracker()
|
||||
tracker.record('inference', 42)
|
||||
self.assertEqual(tracker.total_units, 42)
|
||||
self.assertEqual(len(tracker.events), 1)
|
||||
|
||||
def test_record_multiple_events_accumulates_totals(self) -> None:
|
||||
tracker = CostTracker()
|
||||
tracker.record('inference', 10)
|
||||
tracker.record('embedding', 20)
|
||||
tracker.record('search', 30)
|
||||
self.assertEqual(tracker.total_units, 60)
|
||||
self.assertEqual(len(tracker.events), 3)
|
||||
|
||||
def test_record_zero_units(self) -> None:
|
||||
tracker = CostTracker()
|
||||
tracker.record('noop', 0)
|
||||
self.assertEqual(tracker.total_units, 0)
|
||||
self.assertEqual(len(tracker.events), 1)
|
||||
self.assertIn('noop:0', tracker.events)
|
||||
|
||||
def test_record_large_units(self) -> None:
|
||||
tracker = CostTracker()
|
||||
large = 10**9
|
||||
tracker.record('bulk', large)
|
||||
self.assertEqual(tracker.total_units, large)
|
||||
self.assertEqual(tracker.events, [f'bulk:{large}'])
|
||||
|
||||
def test_event_format_is_label_colon_units(self) -> None:
|
||||
tracker = CostTracker()
|
||||
tracker.record('inference', 42)
|
||||
self.assertEqual(tracker.events[0], 'inference:42')
|
||||
|
||||
def test_events_are_ordered_chronologically(self) -> None:
|
||||
tracker = CostTracker()
|
||||
labels = ['first', 'second', 'third']
|
||||
for i, label in enumerate(labels):
|
||||
tracker.record(label, i)
|
||||
self.assertEqual(tracker.events, ['first:0', 'second:1', 'third:2'])
|
||||
@@ -0,0 +1,179 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.execution_registry import (
|
||||
ExecutionRegistry,
|
||||
MirroredCommand,
|
||||
MirroredTool,
|
||||
build_execution_registry,
|
||||
)
|
||||
from src.commands import PORTED_COMMANDS
|
||||
from src.tools import PORTED_TOOLS
|
||||
|
||||
|
||||
class TestBuildExecutionRegistry(unittest.TestCase):
|
||||
"""Tests for build_execution_registry and basic registry properties."""
|
||||
|
||||
def test_returns_execution_registry(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
self.assertIsInstance(registry, ExecutionRegistry)
|
||||
|
||||
def test_has_non_empty_commands(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
self.assertGreater(len(registry.commands), 0)
|
||||
|
||||
def test_has_non_empty_tools(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
self.assertGreater(len(registry.tools), 0)
|
||||
|
||||
def test_command_count_matches_ported_commands(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
self.assertEqual(len(registry.commands), len(PORTED_COMMANDS))
|
||||
|
||||
def test_tool_count_matches_ported_tools(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
self.assertEqual(len(registry.tools), len(PORTED_TOOLS))
|
||||
|
||||
|
||||
class TestCommandLookup(unittest.TestCase):
|
||||
"""Tests for ExecutionRegistry.command() lookup."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.registry = build_execution_registry()
|
||||
self.known_name = self.registry.commands[0].name
|
||||
|
||||
def test_lookup_exact_case(self) -> None:
|
||||
result = self.registry.command(self.known_name)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_lower(self) -> None:
|
||||
result = self.registry.command(self.known_name.lower())
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_upper(self) -> None:
|
||||
result = self.registry.command(self.known_name.upper())
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_mixed(self) -> None:
|
||||
mixed = ''.join(
|
||||
c.upper() if i % 2 else c.lower()
|
||||
for i, c in enumerate(self.known_name)
|
||||
)
|
||||
result = self.registry.command(mixed)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_returns_none_for_unknown_name(self) -> None:
|
||||
result = self.registry.command('__nonexistent_command_xyz__')
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestToolLookup(unittest.TestCase):
|
||||
"""Tests for ExecutionRegistry.tool() lookup."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.registry = build_execution_registry()
|
||||
self.known_name = self.registry.tools[0].name
|
||||
|
||||
def test_lookup_exact_case(self) -> None:
|
||||
result = self.registry.tool(self.known_name)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_lower(self) -> None:
|
||||
result = self.registry.tool(self.known_name.lower())
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_upper(self) -> None:
|
||||
result = self.registry.tool(self.known_name.upper())
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_lookup_case_insensitive_mixed(self) -> None:
|
||||
mixed = ''.join(
|
||||
c.upper() if i % 2 else c.lower()
|
||||
for i, c in enumerate(self.known_name)
|
||||
)
|
||||
result = self.registry.tool(mixed)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result.name, self.known_name)
|
||||
|
||||
def test_returns_none_for_unknown_name(self) -> None:
|
||||
result = self.registry.tool('__nonexistent_tool_xyz__')
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestMirroredCommand(unittest.TestCase):
|
||||
"""Tests for MirroredCommand dataclass."""
|
||||
|
||||
def test_has_correct_name(self) -> None:
|
||||
cmd = MirroredCommand(name='review', source_hint='copilot')
|
||||
self.assertEqual(cmd.name, 'review')
|
||||
|
||||
def test_has_correct_source_hint(self) -> None:
|
||||
cmd = MirroredCommand(name='review', source_hint='copilot')
|
||||
self.assertEqual(cmd.source_hint, 'copilot')
|
||||
|
||||
def test_is_frozen(self) -> None:
|
||||
cmd = MirroredCommand(name='review', source_hint='copilot')
|
||||
with self.assertRaises(AttributeError):
|
||||
cmd.name = 'other' # type: ignore[misc]
|
||||
|
||||
def test_execute_returns_string(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
cmd = registry.commands[0]
|
||||
result = cmd.execute('test prompt')
|
||||
self.assertIsInstance(result, str)
|
||||
self.assertIn('Mirrored command', result)
|
||||
|
||||
|
||||
class TestMirroredTool(unittest.TestCase):
|
||||
"""Tests for MirroredTool dataclass."""
|
||||
|
||||
def test_has_correct_name(self) -> None:
|
||||
tool = MirroredTool(name='BashTool', source_hint='vscode')
|
||||
self.assertEqual(tool.name, 'BashTool')
|
||||
|
||||
def test_has_correct_source_hint(self) -> None:
|
||||
tool = MirroredTool(name='BashTool', source_hint='vscode')
|
||||
self.assertEqual(tool.source_hint, 'vscode')
|
||||
|
||||
def test_is_frozen(self) -> None:
|
||||
tool = MirroredTool(name='BashTool', source_hint='vscode')
|
||||
with self.assertRaises(AttributeError):
|
||||
tool.name = 'other' # type: ignore[misc]
|
||||
|
||||
def test_execute_returns_string(self) -> None:
|
||||
registry = build_execution_registry()
|
||||
tool = registry.tools[0]
|
||||
result = tool.execute('test payload')
|
||||
self.assertIsInstance(result, str)
|
||||
self.assertIn('Mirrored tool', result)
|
||||
|
||||
|
||||
class TestEmptyRegistry(unittest.TestCase):
|
||||
"""Tests for an empty ExecutionRegistry."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.registry = ExecutionRegistry(commands=(), tools=())
|
||||
|
||||
def test_command_returns_none(self) -> None:
|
||||
self.assertIsNone(self.registry.command('anything'))
|
||||
|
||||
def test_tool_returns_none(self) -> None:
|
||||
self.assertIsNone(self.registry.tool('anything'))
|
||||
|
||||
def test_empty_commands_tuple(self) -> None:
|
||||
self.assertEqual(len(self.registry.commands), 0)
|
||||
|
||||
def test_empty_tools_tuple(self) -> None:
|
||||
self.assertEqual(len(self.registry.tools), 0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,167 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.openai_compat import (
|
||||
OpenAICompatError,
|
||||
_build_response_format,
|
||||
_join_url,
|
||||
_normalize_content,
|
||||
_optional_int,
|
||||
_parse_tool_arguments,
|
||||
_parse_usage,
|
||||
)
|
||||
from src.agent_types import OutputSchemaConfig, UsageStats
|
||||
|
||||
|
||||
class TestJoinUrl(unittest.TestCase):
|
||||
def test_base_with_trailing_slash(self):
|
||||
self.assertEqual(_join_url('http://localhost:8000/', 'v1/chat'), 'http://localhost:8000/v1/chat')
|
||||
|
||||
def test_base_without_trailing_slash(self):
|
||||
self.assertEqual(_join_url('http://localhost:8000', 'v1/chat'), 'http://localhost:8000/v1/chat')
|
||||
|
||||
def test_suffix_with_leading_slash(self):
|
||||
self.assertEqual(_join_url('http://localhost:8000', '/v1/chat'), 'http://localhost:8000/v1/chat')
|
||||
|
||||
|
||||
class TestNormalizeContent(unittest.TestCase):
|
||||
def test_string_passthrough(self):
|
||||
self.assertEqual(_normalize_content('hello'), 'hello')
|
||||
|
||||
def test_none_returns_empty(self):
|
||||
self.assertEqual(_normalize_content(None), '')
|
||||
|
||||
def test_list_of_strings_joined(self):
|
||||
self.assertEqual(_normalize_content(['hello', ' ', 'world']), 'hello world')
|
||||
|
||||
def test_list_of_text_dicts(self):
|
||||
items = [{'type': 'text', 'text': 'hello'}, {'type': 'text', 'text': ' world'}]
|
||||
self.assertEqual(_normalize_content(items), 'hello world')
|
||||
|
||||
def test_list_of_mixed_items(self):
|
||||
items = ['start ', {'type': 'text', 'text': 'middle'}, ' end']
|
||||
self.assertEqual(_normalize_content(items), 'start middle end')
|
||||
|
||||
def test_non_string_non_list_returns_str(self):
|
||||
self.assertEqual(_normalize_content(42), '42')
|
||||
|
||||
|
||||
class TestParseToolArguments(unittest.TestCase):
|
||||
def test_dict_passthrough(self):
|
||||
d = {'key': 'value'}
|
||||
self.assertIs(_parse_tool_arguments(d), d)
|
||||
|
||||
def test_valid_json_string(self):
|
||||
self.assertEqual(_parse_tool_arguments('{"a": 1}'), {'a': 1})
|
||||
|
||||
def test_empty_string_returns_empty_dict(self):
|
||||
self.assertEqual(_parse_tool_arguments(''), {})
|
||||
|
||||
def test_none_returns_empty_dict(self):
|
||||
self.assertEqual(_parse_tool_arguments(None), {})
|
||||
|
||||
def test_invalid_json_raises(self):
|
||||
with self.assertRaises(OpenAICompatError):
|
||||
_parse_tool_arguments('{bad json}')
|
||||
|
||||
def test_json_non_dict_raises(self):
|
||||
with self.assertRaises(OpenAICompatError):
|
||||
_parse_tool_arguments('[1, 2, 3]')
|
||||
|
||||
def test_unsupported_type_raises(self):
|
||||
with self.assertRaises(OpenAICompatError):
|
||||
_parse_tool_arguments(12345)
|
||||
|
||||
|
||||
class TestParseUsage(unittest.TestCase):
|
||||
def test_standard_fields(self):
|
||||
usage = _parse_usage({'input_tokens': 10, 'output_tokens': 20})
|
||||
self.assertEqual(usage.input_tokens, 10)
|
||||
self.assertEqual(usage.output_tokens, 20)
|
||||
|
||||
def test_prompt_completion_aliases(self):
|
||||
usage = _parse_usage({'prompt_tokens': 15, 'completion_tokens': 25})
|
||||
self.assertEqual(usage.input_tokens, 15)
|
||||
self.assertEqual(usage.output_tokens, 25)
|
||||
|
||||
def test_ollama_aliases(self):
|
||||
usage = _parse_usage({'prompt_eval_count': 12, 'eval_count': 18})
|
||||
self.assertEqual(usage.input_tokens, 12)
|
||||
self.assertEqual(usage.output_tokens, 18)
|
||||
|
||||
def test_cache_tokens(self):
|
||||
usage = _parse_usage({
|
||||
'input_tokens': 1,
|
||||
'output_tokens': 1,
|
||||
'cache_creation_input_tokens': 100,
|
||||
'cache_read_input_tokens': 200,
|
||||
})
|
||||
self.assertEqual(usage.cache_creation_input_tokens, 100)
|
||||
self.assertEqual(usage.cache_read_input_tokens, 200)
|
||||
|
||||
def test_reasoning_tokens_top_level_and_details(self):
|
||||
usage_top = _parse_usage({'input_tokens': 1, 'output_tokens': 1, 'reasoning_tokens': 50})
|
||||
self.assertEqual(usage_top.reasoning_tokens, 50)
|
||||
|
||||
usage_details = _parse_usage({
|
||||
'input_tokens': 1,
|
||||
'output_tokens': 1,
|
||||
'completion_tokens_details': {'reasoning_tokens': 75},
|
||||
})
|
||||
self.assertEqual(usage_details.reasoning_tokens, 75)
|
||||
|
||||
def test_non_dict_returns_empty(self):
|
||||
usage = _parse_usage('not a dict')
|
||||
self.assertEqual(usage, UsageStats())
|
||||
|
||||
def test_string_number_coercion(self):
|
||||
usage = _parse_usage({'input_tokens': '10', 'output_tokens': '20'})
|
||||
self.assertEqual(usage.input_tokens, 10)
|
||||
self.assertEqual(usage.output_tokens, 20)
|
||||
|
||||
|
||||
class TestBuildResponseFormat(unittest.TestCase):
|
||||
def test_none_returns_none(self):
|
||||
self.assertIsNone(_build_response_format(None))
|
||||
|
||||
def test_valid_schema(self):
|
||||
schema = OutputSchemaConfig(
|
||||
name='test_schema',
|
||||
schema={'type': 'object', 'properties': {'x': {'type': 'integer'}}},
|
||||
strict=True,
|
||||
)
|
||||
result = _build_response_format(schema)
|
||||
self.assertEqual(result, {
|
||||
'type': 'json_schema',
|
||||
'json_schema': {
|
||||
'name': 'test_schema',
|
||||
'schema': {'type': 'object', 'properties': {'x': {'type': 'integer'}}},
|
||||
'strict': True,
|
||||
},
|
||||
})
|
||||
|
||||
|
||||
class TestOptionalInt(unittest.TestCase):
|
||||
def test_int_passthrough(self):
|
||||
self.assertEqual(_optional_int(42), 42)
|
||||
|
||||
def test_float_truncated(self):
|
||||
self.assertEqual(_optional_int(3.9), 3)
|
||||
|
||||
def test_string_parsed(self):
|
||||
self.assertEqual(_optional_int('7'), 7)
|
||||
|
||||
def test_bool_returns_zero(self):
|
||||
self.assertEqual(_optional_int(True), 0)
|
||||
self.assertEqual(_optional_int(False), 0)
|
||||
|
||||
def test_none_returns_zero(self):
|
||||
self.assertEqual(_optional_int(None), 0)
|
||||
|
||||
def test_invalid_string_returns_zero(self):
|
||||
self.assertEqual(_optional_int('abc'), 0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,87 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.permissions import ToolPermissionContext
|
||||
|
||||
|
||||
class TestToolPermissionContext(unittest.TestCase):
|
||||
# 1. Empty context blocks nothing
|
||||
def test_empty_context_blocks_nothing(self) -> None:
|
||||
ctx = ToolPermissionContext()
|
||||
self.assertFalse(ctx.blocks("anything"))
|
||||
self.assertFalse(ctx.blocks(""))
|
||||
|
||||
# 2. Exact name blocking (case insensitive)
|
||||
def test_exact_name_blocking(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(deny_names=["dangerous_tool"])
|
||||
self.assertTrue(ctx.blocks("dangerous_tool"))
|
||||
self.assertTrue(ctx.blocks("Dangerous_Tool"))
|
||||
self.assertTrue(ctx.blocks("DANGEROUS_TOOL"))
|
||||
self.assertFalse(ctx.blocks("safe_tool"))
|
||||
|
||||
# 3. Prefix blocking (case insensitive)
|
||||
def test_prefix_blocking(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(deny_prefixes=["admin_"])
|
||||
self.assertTrue(ctx.blocks("admin_delete"))
|
||||
self.assertTrue(ctx.blocks("Admin_Delete"))
|
||||
self.assertTrue(ctx.blocks("ADMIN_CREATE"))
|
||||
self.assertFalse(ctx.blocks("user_admin"))
|
||||
|
||||
# 4. Combined name + prefix blocking
|
||||
def test_combined_name_and_prefix_blocking(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(
|
||||
deny_names=["rm"],
|
||||
deny_prefixes=["sudo_"],
|
||||
)
|
||||
self.assertTrue(ctx.blocks("rm"))
|
||||
self.assertTrue(ctx.blocks("sudo_restart"))
|
||||
self.assertFalse(ctx.blocks("ls"))
|
||||
|
||||
# 5. Non-matching names are allowed
|
||||
def test_non_matching_names_allowed(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(
|
||||
deny_names=["blocked"],
|
||||
deny_prefixes=["bad_"],
|
||||
)
|
||||
self.assertFalse(ctx.blocks("allowed"))
|
||||
self.assertFalse(ctx.blocks("good_tool"))
|
||||
self.assertFalse(ctx.blocks("not_bad"))
|
||||
|
||||
# 6. from_iterables with None args
|
||||
def test_from_iterables_none_args(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(deny_names=None, deny_prefixes=None)
|
||||
self.assertEqual(ctx.deny_names, frozenset())
|
||||
self.assertEqual(ctx.deny_prefixes, ())
|
||||
self.assertFalse(ctx.blocks("anything"))
|
||||
|
||||
def test_from_iterables_default_args(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables()
|
||||
self.assertEqual(ctx.deny_names, frozenset())
|
||||
self.assertEqual(ctx.deny_prefixes, ())
|
||||
|
||||
# 7. from_iterables normalizes to lowercase
|
||||
def test_from_iterables_normalizes_to_lowercase(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(
|
||||
deny_names=["FooBar"],
|
||||
deny_prefixes=["PFX_"],
|
||||
)
|
||||
self.assertIn("foobar", ctx.deny_names)
|
||||
self.assertNotIn("FooBar", ctx.deny_names)
|
||||
self.assertEqual(ctx.deny_prefixes, ("pfx_",))
|
||||
self.assertTrue(ctx.blocks("FOOBAR"))
|
||||
self.assertTrue(ctx.blocks("pfx_something"))
|
||||
|
||||
# 8. Multiple deny_names
|
||||
def test_multiple_deny_names(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(
|
||||
deny_names=["tool_a", "tool_b", "tool_c"],
|
||||
)
|
||||
self.assertTrue(ctx.blocks("tool_a"))
|
||||
self.assertTrue(ctx.blocks("tool_b"))
|
||||
self.assertTrue(ctx.blocks("tool_c"))
|
||||
self.assertFalse(ctx.blocks("tool_d"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,553 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from src.agent_types import (
|
||||
AgentPermissions,
|
||||
AgentRuntimeConfig,
|
||||
BudgetConfig,
|
||||
ModelConfig,
|
||||
ModelPricing,
|
||||
OutputSchemaConfig,
|
||||
UsageStats,
|
||||
)
|
||||
from src.session_store import (
|
||||
StoredAgentSession,
|
||||
StoredSession,
|
||||
_deserialize_output_schema,
|
||||
_optional_float,
|
||||
_optional_int,
|
||||
deserialize_model_config,
|
||||
deserialize_runtime_config,
|
||||
load_agent_session,
|
||||
load_session,
|
||||
save_agent_session,
|
||||
save_session,
|
||||
serialize_model_config,
|
||||
serialize_runtime_config,
|
||||
usage_from_payload,
|
||||
)
|
||||
|
||||
|
||||
class TestStoredSessionRoundTrip(unittest.TestCase):
|
||||
"""save_session then load_session preserves all fields."""
|
||||
|
||||
def test_round_trip(self) -> None:
|
||||
session = StoredSession(
|
||||
session_id='abc-123',
|
||||
messages=('hello', 'world', 'foo'),
|
||||
input_tokens=100,
|
||||
output_tokens=200,
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
save_session(session, directory=directory)
|
||||
loaded = load_session('abc-123', directory=directory)
|
||||
|
||||
self.assertEqual(loaded.session_id, session.session_id)
|
||||
self.assertEqual(loaded.messages, session.messages)
|
||||
self.assertEqual(loaded.input_tokens, session.input_tokens)
|
||||
self.assertEqual(loaded.output_tokens, session.output_tokens)
|
||||
|
||||
def test_round_trip_empty_messages(self) -> None:
|
||||
session = StoredSession(
|
||||
session_id='empty',
|
||||
messages=(),
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
save_session(session, directory=directory)
|
||||
loaded = load_session('empty', directory=directory)
|
||||
|
||||
self.assertEqual(loaded.messages, ())
|
||||
self.assertEqual(loaded.input_tokens, 0)
|
||||
|
||||
|
||||
class TestStoredAgentSessionRoundTrip(unittest.TestCase):
|
||||
"""save_agent_session then load_agent_session preserves all fields."""
|
||||
|
||||
def _make_session(self, **overrides: object) -> StoredAgentSession:
|
||||
defaults: dict = {
|
||||
'session_id': 'agent-001',
|
||||
'model_config': {'model': 'gpt-4', 'temperature': 0.5},
|
||||
'runtime_config': {'cwd': '/home/user', 'max_turns': 20},
|
||||
'system_prompt_parts': ('You are helpful.',),
|
||||
'user_context': {'lang': 'en'},
|
||||
'system_context': {'os': 'linux'},
|
||||
'messages': ({'role': 'user', 'content': 'hi'},),
|
||||
'turns': 3,
|
||||
'tool_calls': 7,
|
||||
'usage': {'input_tokens': 500, 'output_tokens': 300},
|
||||
'total_cost_usd': 0.05,
|
||||
'file_history': ({'file': 'a.py', 'action': 'edit'},),
|
||||
'budget_state': {'remaining': 100},
|
||||
'plugin_state': {'key': 'value'},
|
||||
'scratchpad_directory': '/scratch/pad',
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return StoredAgentSession(**defaults)
|
||||
|
||||
def test_round_trip_all_fields(self) -> None:
|
||||
session = self._make_session()
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
save_agent_session(session, directory=directory)
|
||||
loaded = load_agent_session('agent-001', directory=directory)
|
||||
|
||||
self.assertEqual(loaded.session_id, session.session_id)
|
||||
self.assertEqual(loaded.model_config, session.model_config)
|
||||
self.assertEqual(loaded.runtime_config, session.runtime_config)
|
||||
self.assertEqual(loaded.system_prompt_parts, session.system_prompt_parts)
|
||||
self.assertEqual(loaded.user_context, session.user_context)
|
||||
self.assertEqual(loaded.system_context, session.system_context)
|
||||
self.assertEqual(loaded.messages, session.messages)
|
||||
self.assertEqual(loaded.turns, session.turns)
|
||||
self.assertEqual(loaded.tool_calls, session.tool_calls)
|
||||
self.assertEqual(loaded.usage, session.usage)
|
||||
self.assertAlmostEqual(loaded.total_cost_usd, session.total_cost_usd)
|
||||
self.assertEqual(loaded.file_history, session.file_history)
|
||||
self.assertEqual(loaded.budget_state, session.budget_state)
|
||||
self.assertEqual(loaded.plugin_state, session.plugin_state)
|
||||
self.assertEqual(loaded.scratchpad_directory, session.scratchpad_directory)
|
||||
|
||||
def test_round_trip_no_scratchpad(self) -> None:
|
||||
session = self._make_session(scratchpad_directory=None)
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
save_agent_session(session, directory=directory)
|
||||
loaded = load_agent_session('agent-001', directory=directory)
|
||||
|
||||
self.assertIsNone(loaded.scratchpad_directory)
|
||||
|
||||
def test_load_filters_non_dict_messages(self) -> None:
|
||||
"""Non-dict entries in messages list are filtered out on load."""
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
path = directory / 'mixed.json'
|
||||
data = {
|
||||
'session_id': 'mixed',
|
||||
'model_config': {},
|
||||
'runtime_config': {'cwd': '/'},
|
||||
'system_prompt_parts': [],
|
||||
'user_context': {},
|
||||
'system_context': {},
|
||||
'messages': [
|
||||
{'role': 'user', 'content': 'hi'},
|
||||
'not a dict',
|
||||
42,
|
||||
None,
|
||||
{'role': 'assistant', 'content': 'hey'},
|
||||
],
|
||||
'turns': 0,
|
||||
'tool_calls': 0,
|
||||
'usage': {},
|
||||
'total_cost_usd': 0.0,
|
||||
'file_history': [],
|
||||
'budget_state': {},
|
||||
'plugin_state': {},
|
||||
}
|
||||
path.write_text(json.dumps(data))
|
||||
loaded = load_agent_session('mixed', directory=directory)
|
||||
|
||||
self.assertEqual(len(loaded.messages), 2)
|
||||
self.assertEqual(loaded.messages[0]['role'], 'user')
|
||||
self.assertEqual(loaded.messages[1]['role'], 'assistant')
|
||||
|
||||
def test_load_defaults_for_missing_optional_fields(self) -> None:
|
||||
"""Missing optional fields get sensible defaults."""
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
path = directory / 'minimal.json'
|
||||
data = {
|
||||
'session_id': 'minimal',
|
||||
'model_config': {},
|
||||
'runtime_config': {'cwd': '/'},
|
||||
'system_prompt_parts': [],
|
||||
'user_context': {},
|
||||
'system_context': {},
|
||||
'messages': [],
|
||||
'turns': 1,
|
||||
'tool_calls': 2,
|
||||
}
|
||||
path.write_text(json.dumps(data))
|
||||
loaded = load_agent_session('minimal', directory=directory)
|
||||
|
||||
self.assertEqual(loaded.usage, {})
|
||||
self.assertAlmostEqual(loaded.total_cost_usd, 0.0)
|
||||
self.assertEqual(loaded.file_history, ())
|
||||
self.assertEqual(loaded.budget_state, {})
|
||||
self.assertEqual(loaded.plugin_state, {})
|
||||
self.assertIsNone(loaded.scratchpad_directory)
|
||||
|
||||
def test_load_non_dict_budget_state_defaults_to_empty(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
directory = Path(td)
|
||||
path = directory / 'bad-budget.json'
|
||||
data = {
|
||||
'session_id': 'bad-budget',
|
||||
'model_config': {},
|
||||
'runtime_config': {'cwd': '/'},
|
||||
'system_prompt_parts': [],
|
||||
'user_context': {},
|
||||
'system_context': {},
|
||||
'messages': [],
|
||||
'turns': 0,
|
||||
'tool_calls': 0,
|
||||
'budget_state': 'not-a-dict',
|
||||
'plugin_state': 123,
|
||||
}
|
||||
path.write_text(json.dumps(data))
|
||||
loaded = load_agent_session('bad-budget', directory=directory)
|
||||
|
||||
self.assertEqual(loaded.budget_state, {})
|
||||
self.assertEqual(loaded.plugin_state, {})
|
||||
|
||||
|
||||
class TestModelConfigSerialization(unittest.TestCase):
|
||||
"""serialize_model_config + deserialize_model_config round-trip."""
|
||||
|
||||
def test_round_trip_preserves_pricing(self) -> None:
|
||||
pricing = ModelPricing(
|
||||
input_cost_per_million_tokens_usd=3.0,
|
||||
output_cost_per_million_tokens_usd=15.0,
|
||||
cache_creation_input_cost_per_million_tokens_usd=1.5,
|
||||
cache_read_input_cost_per_million_tokens_usd=0.5,
|
||||
)
|
||||
config = ModelConfig(
|
||||
model='claude-3-sonnet',
|
||||
base_url='https://api.example.com/v1',
|
||||
api_key='sk-test-key',
|
||||
temperature=0.7,
|
||||
timeout_seconds=60.0,
|
||||
pricing=pricing,
|
||||
)
|
||||
payload = serialize_model_config(config)
|
||||
restored = deserialize_model_config(payload)
|
||||
|
||||
self.assertEqual(restored.model, config.model)
|
||||
self.assertEqual(restored.base_url, config.base_url)
|
||||
self.assertEqual(restored.api_key, config.api_key)
|
||||
self.assertAlmostEqual(restored.temperature, config.temperature)
|
||||
self.assertAlmostEqual(restored.timeout_seconds, config.timeout_seconds)
|
||||
self.assertAlmostEqual(
|
||||
restored.pricing.input_cost_per_million_tokens_usd,
|
||||
pricing.input_cost_per_million_tokens_usd,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
restored.pricing.output_cost_per_million_tokens_usd,
|
||||
pricing.output_cost_per_million_tokens_usd,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
restored.pricing.cache_creation_input_cost_per_million_tokens_usd,
|
||||
pricing.cache_creation_input_cost_per_million_tokens_usd,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
restored.pricing.cache_read_input_cost_per_million_tokens_usd,
|
||||
pricing.cache_read_input_cost_per_million_tokens_usd,
|
||||
)
|
||||
|
||||
def test_deserialize_defaults_for_missing_fields(self) -> None:
|
||||
payload = {'model': 'gpt-4'}
|
||||
config = deserialize_model_config(payload)
|
||||
|
||||
self.assertEqual(config.model, 'gpt-4')
|
||||
self.assertEqual(config.base_url, 'http://127.0.0.1:8000/v1')
|
||||
self.assertEqual(config.api_key, 'local-token')
|
||||
self.assertAlmostEqual(config.temperature, 0.0)
|
||||
self.assertAlmostEqual(config.timeout_seconds, 120.0)
|
||||
self.assertAlmostEqual(config.pricing.input_cost_per_million_tokens_usd, 0.0)
|
||||
self.assertAlmostEqual(config.pricing.output_cost_per_million_tokens_usd, 0.0)
|
||||
|
||||
def test_deserialize_with_non_dict_pricing(self) -> None:
|
||||
payload = {'model': 'test', 'pricing': 'invalid'}
|
||||
config = deserialize_model_config(payload)
|
||||
self.assertAlmostEqual(config.pricing.input_cost_per_million_tokens_usd, 0.0)
|
||||
|
||||
def test_deserialize_with_none_pricing(self) -> None:
|
||||
payload = {'model': 'test', 'pricing': None}
|
||||
config = deserialize_model_config(payload)
|
||||
self.assertEqual(config.pricing, ModelPricing())
|
||||
|
||||
|
||||
class TestRuntimeConfigSerialization(unittest.TestCase):
|
||||
"""serialize_runtime_config + deserialize_runtime_config round-trip."""
|
||||
|
||||
def test_round_trip_preserves_all(self) -> None:
|
||||
config = AgentRuntimeConfig(
|
||||
cwd=Path('/home/user/project'),
|
||||
max_turns=25,
|
||||
command_timeout_seconds=45.0,
|
||||
max_output_chars=8000,
|
||||
stream_model_responses=True,
|
||||
auto_snip_threshold_tokens=5000,
|
||||
auto_compact_threshold_tokens=10000,
|
||||
compact_preserve_messages=6,
|
||||
permissions=AgentPermissions(
|
||||
allow_file_write=True,
|
||||
allow_shell_commands=True,
|
||||
allow_destructive_shell_commands=False,
|
||||
),
|
||||
additional_working_directories=(Path('/extra/dir'),),
|
||||
disable_claude_md_discovery=True,
|
||||
budget_config=BudgetConfig(
|
||||
max_total_tokens=100000,
|
||||
max_input_tokens=50000,
|
||||
max_output_tokens=30000,
|
||||
max_reasoning_tokens=20000,
|
||||
max_total_cost_usd=5.0,
|
||||
max_tool_calls=100,
|
||||
max_delegated_tasks=10,
|
||||
max_model_calls=200,
|
||||
max_session_turns=50,
|
||||
),
|
||||
output_schema=OutputSchemaConfig(
|
||||
name='test_schema',
|
||||
schema={'type': 'object', 'properties': {'answer': {'type': 'string'}}},
|
||||
strict=True,
|
||||
),
|
||||
session_directory=Path('/sessions'),
|
||||
scratchpad_root=Path('/scratch'),
|
||||
)
|
||||
payload = serialize_runtime_config(config)
|
||||
restored = deserialize_runtime_config(payload)
|
||||
|
||||
self.assertEqual(restored.cwd, config.cwd.resolve())
|
||||
self.assertEqual(restored.max_turns, 25)
|
||||
self.assertAlmostEqual(restored.command_timeout_seconds, 45.0)
|
||||
self.assertEqual(restored.max_output_chars, 8000)
|
||||
self.assertTrue(restored.stream_model_responses)
|
||||
self.assertEqual(restored.auto_snip_threshold_tokens, 5000)
|
||||
self.assertEqual(restored.auto_compact_threshold_tokens, 10000)
|
||||
self.assertEqual(restored.compact_preserve_messages, 6)
|
||||
self.assertTrue(restored.permissions.allow_file_write)
|
||||
self.assertTrue(restored.permissions.allow_shell_commands)
|
||||
self.assertFalse(restored.permissions.allow_destructive_shell_commands)
|
||||
self.assertTrue(restored.disable_claude_md_discovery)
|
||||
|
||||
self.assertEqual(restored.budget_config.max_total_tokens, 100000)
|
||||
self.assertEqual(restored.budget_config.max_input_tokens, 50000)
|
||||
self.assertEqual(restored.budget_config.max_output_tokens, 30000)
|
||||
self.assertEqual(restored.budget_config.max_reasoning_tokens, 20000)
|
||||
self.assertAlmostEqual(restored.budget_config.max_total_cost_usd, 5.0)
|
||||
self.assertEqual(restored.budget_config.max_tool_calls, 100)
|
||||
self.assertEqual(restored.budget_config.max_delegated_tasks, 10)
|
||||
self.assertEqual(restored.budget_config.max_model_calls, 200)
|
||||
self.assertEqual(restored.budget_config.max_session_turns, 50)
|
||||
|
||||
self.assertIsNotNone(restored.output_schema)
|
||||
assert restored.output_schema is not None
|
||||
self.assertEqual(restored.output_schema.name, 'test_schema')
|
||||
self.assertEqual(restored.output_schema.schema, config.output_schema.schema)
|
||||
self.assertTrue(restored.output_schema.strict)
|
||||
|
||||
def test_round_trip_none_output_schema(self) -> None:
|
||||
config = AgentRuntimeConfig(
|
||||
cwd=Path('/home/user'),
|
||||
output_schema=None,
|
||||
)
|
||||
payload = serialize_runtime_config(config)
|
||||
restored = deserialize_runtime_config(payload)
|
||||
self.assertIsNone(restored.output_schema)
|
||||
|
||||
def test_deserialize_defaults_for_missing_fields(self) -> None:
|
||||
payload = {'cwd': '/home/user'}
|
||||
config = deserialize_runtime_config(payload)
|
||||
|
||||
self.assertEqual(config.max_turns, 12)
|
||||
self.assertAlmostEqual(config.command_timeout_seconds, 30.0)
|
||||
self.assertEqual(config.max_output_chars, 12000)
|
||||
self.assertFalse(config.stream_model_responses)
|
||||
self.assertIsNone(config.auto_snip_threshold_tokens)
|
||||
self.assertIsNone(config.auto_compact_threshold_tokens)
|
||||
self.assertEqual(config.compact_preserve_messages, 4)
|
||||
self.assertFalse(config.permissions.allow_file_write)
|
||||
self.assertFalse(config.permissions.allow_shell_commands)
|
||||
self.assertFalse(config.permissions.allow_destructive_shell_commands)
|
||||
self.assertEqual(config.additional_working_directories, ())
|
||||
self.assertFalse(config.disable_claude_md_discovery)
|
||||
self.assertIsNone(config.budget_config.max_total_tokens)
|
||||
self.assertIsNone(config.output_schema)
|
||||
|
||||
def test_deserialize_non_dict_permissions(self) -> None:
|
||||
payload = {'cwd': '/home', 'permissions': 'invalid'}
|
||||
config = deserialize_runtime_config(payload)
|
||||
self.assertFalse(config.permissions.allow_file_write)
|
||||
|
||||
def test_deserialize_non_dict_budget_config(self) -> None:
|
||||
payload = {'cwd': '/home', 'budget_config': 42}
|
||||
config = deserialize_runtime_config(payload)
|
||||
self.assertIsNone(config.budget_config.max_total_tokens)
|
||||
|
||||
|
||||
class TestUsageFromPayload(unittest.TestCase):
|
||||
"""usage_from_payload correctly maps fields including defaults."""
|
||||
|
||||
def test_full_payload(self) -> None:
|
||||
payload = {
|
||||
'input_tokens': 1000,
|
||||
'output_tokens': 500,
|
||||
'cache_creation_input_tokens': 200,
|
||||
'cache_read_input_tokens': 100,
|
||||
'reasoning_tokens': 50,
|
||||
}
|
||||
usage = usage_from_payload(payload)
|
||||
self.assertEqual(usage.input_tokens, 1000)
|
||||
self.assertEqual(usage.output_tokens, 500)
|
||||
self.assertEqual(usage.cache_creation_input_tokens, 200)
|
||||
self.assertEqual(usage.cache_read_input_tokens, 100)
|
||||
self.assertEqual(usage.reasoning_tokens, 50)
|
||||
|
||||
def test_partial_payload_uses_defaults(self) -> None:
|
||||
payload = {'input_tokens': 42}
|
||||
usage = usage_from_payload(payload)
|
||||
self.assertEqual(usage.input_tokens, 42)
|
||||
self.assertEqual(usage.output_tokens, 0)
|
||||
self.assertEqual(usage.cache_creation_input_tokens, 0)
|
||||
self.assertEqual(usage.cache_read_input_tokens, 0)
|
||||
self.assertEqual(usage.reasoning_tokens, 0)
|
||||
|
||||
def test_none_returns_empty(self) -> None:
|
||||
usage = usage_from_payload(None)
|
||||
self.assertEqual(usage, UsageStats())
|
||||
|
||||
def test_empty_dict_returns_defaults(self) -> None:
|
||||
usage = usage_from_payload({})
|
||||
self.assertEqual(usage, UsageStats())
|
||||
|
||||
def test_non_dict_returns_empty(self) -> None:
|
||||
usage = usage_from_payload('not a dict') # type: ignore[arg-type]
|
||||
self.assertEqual(usage, UsageStats())
|
||||
|
||||
def test_string_token_values_parsed(self) -> None:
|
||||
payload = {'input_tokens': '99', 'output_tokens': '77'}
|
||||
usage = usage_from_payload(payload)
|
||||
self.assertEqual(usage.input_tokens, 99)
|
||||
self.assertEqual(usage.output_tokens, 77)
|
||||
|
||||
|
||||
class TestOptionalInt(unittest.TestCase):
|
||||
"""_optional_int handles int, str, float, None, bool correctly."""
|
||||
|
||||
def test_int_value(self) -> None:
|
||||
self.assertEqual(_optional_int(42), 42)
|
||||
|
||||
def test_zero(self) -> None:
|
||||
self.assertEqual(_optional_int(0), 0)
|
||||
|
||||
def test_negative(self) -> None:
|
||||
self.assertEqual(_optional_int(-5), -5)
|
||||
|
||||
def test_str_numeric(self) -> None:
|
||||
self.assertEqual(_optional_int('123'), 123)
|
||||
|
||||
def test_float_value(self) -> None:
|
||||
self.assertEqual(_optional_int(3.9), 3)
|
||||
|
||||
def test_none_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_int(None))
|
||||
|
||||
def test_bool_true_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_int(True))
|
||||
|
||||
def test_bool_false_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_int(False))
|
||||
|
||||
def test_non_numeric_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_int('hello'))
|
||||
|
||||
def test_empty_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_int(''))
|
||||
|
||||
|
||||
class TestOptionalFloat(unittest.TestCase):
|
||||
"""_optional_float handles int, str, float, None, bool correctly."""
|
||||
|
||||
def test_float_value(self) -> None:
|
||||
self.assertAlmostEqual(_optional_float(3.14), 3.14)
|
||||
|
||||
def test_int_value(self) -> None:
|
||||
self.assertAlmostEqual(_optional_float(42), 42.0)
|
||||
|
||||
def test_zero(self) -> None:
|
||||
self.assertAlmostEqual(_optional_float(0), 0.0)
|
||||
|
||||
def test_str_numeric(self) -> None:
|
||||
self.assertAlmostEqual(_optional_float('2.5'), 2.5)
|
||||
|
||||
def test_none_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_float(None))
|
||||
|
||||
def test_bool_true_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_float(True))
|
||||
|
||||
def test_bool_false_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_float(False))
|
||||
|
||||
def test_non_numeric_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_float('abc'))
|
||||
|
||||
def test_empty_string_returns_none(self) -> None:
|
||||
self.assertIsNone(_optional_float(''))
|
||||
|
||||
|
||||
class TestDeserializeOutputSchema(unittest.TestCase):
|
||||
"""_deserialize_output_schema with valid, None, invalid data."""
|
||||
|
||||
def test_valid_payload(self) -> None:
|
||||
payload = {
|
||||
'name': 'my_schema',
|
||||
'schema': {'type': 'object'},
|
||||
'strict': True,
|
||||
}
|
||||
result = _deserialize_output_schema(payload)
|
||||
self.assertIsNotNone(result)
|
||||
assert result is not None
|
||||
self.assertEqual(result.name, 'my_schema')
|
||||
self.assertEqual(result.schema, {'type': 'object'})
|
||||
self.assertTrue(result.strict)
|
||||
|
||||
def test_strict_defaults_false(self) -> None:
|
||||
payload = {
|
||||
'name': 'basic',
|
||||
'schema': {'type': 'string'},
|
||||
}
|
||||
result = _deserialize_output_schema(payload)
|
||||
self.assertIsNotNone(result)
|
||||
assert result is not None
|
||||
self.assertFalse(result.strict)
|
||||
|
||||
def test_none_payload(self) -> None:
|
||||
self.assertIsNone(_deserialize_output_schema(None))
|
||||
|
||||
def test_non_dict_payload(self) -> None:
|
||||
self.assertIsNone(_deserialize_output_schema('not a dict'))
|
||||
self.assertIsNone(_deserialize_output_schema(42))
|
||||
self.assertIsNone(_deserialize_output_schema([]))
|
||||
|
||||
def test_missing_schema_key(self) -> None:
|
||||
self.assertIsNone(_deserialize_output_schema({'name': 'test'}))
|
||||
|
||||
def test_non_dict_schema(self) -> None:
|
||||
self.assertIsNone(_deserialize_output_schema({'name': 'test', 'schema': 'bad'}))
|
||||
|
||||
def test_missing_name(self) -> None:
|
||||
self.assertIsNone(_deserialize_output_schema({'schema': {'type': 'object'}}))
|
||||
|
||||
def test_empty_name(self) -> None:
|
||||
self.assertIsNone(
|
||||
_deserialize_output_schema({'name': '', 'schema': {'type': 'object'}})
|
||||
)
|
||||
|
||||
def test_non_string_name(self) -> None:
|
||||
self.assertIsNone(
|
||||
_deserialize_output_schema({'name': 123, 'schema': {'type': 'object'}})
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,106 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from src.models import PortingModule
|
||||
from src.permissions import ToolPermissionContext
|
||||
from src.tool_pool import ToolPool, assemble_tool_pool
|
||||
|
||||
|
||||
class TestAssembleToolPool(unittest.TestCase):
|
||||
def test_returns_tool_pool_with_tools(self) -> None:
|
||||
pool = assemble_tool_pool()
|
||||
self.assertIsInstance(pool, ToolPool)
|
||||
self.assertIsInstance(pool.tools, tuple)
|
||||
self.assertTrue(all(isinstance(t, PortingModule) for t in pool.tools))
|
||||
|
||||
def test_default_mode_includes_tools(self) -> None:
|
||||
pool = assemble_tool_pool()
|
||||
self.assertGreater(len(pool.tools), 0)
|
||||
|
||||
def test_simple_mode_flag_stored(self) -> None:
|
||||
pool_default = assemble_tool_pool()
|
||||
pool_simple = assemble_tool_pool(simple_mode=True)
|
||||
self.assertFalse(pool_default.simple_mode)
|
||||
self.assertTrue(pool_simple.simple_mode)
|
||||
|
||||
def test_include_mcp_flag_stored(self) -> None:
|
||||
pool_default = assemble_tool_pool()
|
||||
pool_no_mcp = assemble_tool_pool(include_mcp=False)
|
||||
self.assertTrue(pool_default.include_mcp)
|
||||
self.assertFalse(pool_no_mcp.include_mcp)
|
||||
|
||||
def test_simple_mode_reduces_tools(self) -> None:
|
||||
pool_full = assemble_tool_pool(simple_mode=False)
|
||||
pool_simple = assemble_tool_pool(simple_mode=True)
|
||||
self.assertGreater(len(pool_full.tools), len(pool_simple.tools))
|
||||
simple_names = {t.name for t in pool_simple.tools}
|
||||
self.assertTrue(simple_names.issubset({'BashTool', 'FileReadTool', 'FileEditTool'}))
|
||||
|
||||
def test_include_mcp_false_excludes_mcp_tools(self) -> None:
|
||||
pool = assemble_tool_pool(include_mcp=False)
|
||||
for tool in pool.tools:
|
||||
self.assertNotIn('mcp', tool.name.lower())
|
||||
self.assertNotIn('mcp', tool.source_hint.lower())
|
||||
|
||||
def test_permission_context_filters_blocked_tools(self) -> None:
|
||||
ctx = ToolPermissionContext.from_iterables(deny_names=['BashTool'])
|
||||
pool = assemble_tool_pool(permission_context=ctx)
|
||||
tool_names = {t.name for t in pool.tools}
|
||||
self.assertNotIn('BashTool', tool_names)
|
||||
|
||||
pool_unfiltered = assemble_tool_pool()
|
||||
unfiltered_names = {t.name for t in pool_unfiltered.tools}
|
||||
self.assertIn('BashTool', unfiltered_names)
|
||||
|
||||
|
||||
class TestToolPoolAsMarkdown(unittest.TestCase):
|
||||
def test_includes_header_and_tool_count(self) -> None:
|
||||
pool = assemble_tool_pool()
|
||||
md = pool.as_markdown()
|
||||
self.assertIn('# Tool Pool', md)
|
||||
self.assertIn(f'Tool count: {len(pool.tools)}', md)
|
||||
|
||||
def test_includes_mode_flags(self) -> None:
|
||||
pool = assemble_tool_pool(simple_mode=True, include_mcp=False)
|
||||
md = pool.as_markdown()
|
||||
self.assertIn('Simple mode: True', md)
|
||||
self.assertIn('Include MCP: False', md)
|
||||
|
||||
def test_shows_at_most_15_tools(self) -> None:
|
||||
pool = assemble_tool_pool()
|
||||
self.assertGreater(len(pool.tools), 15, 'Need >15 tools for this test')
|
||||
md = pool.as_markdown()
|
||||
tool_lines = [line for line in md.splitlines() if line.startswith('- ')]
|
||||
self.assertEqual(len(tool_lines), 15)
|
||||
|
||||
def test_empty_tools_renders_correctly(self) -> None:
|
||||
pool = ToolPool(tools=(), simple_mode=False, include_mcp=True)
|
||||
md = pool.as_markdown()
|
||||
self.assertIn('# Tool Pool', md)
|
||||
self.assertIn('Tool count: 0', md)
|
||||
self.assertIn('Simple mode: False', md)
|
||||
self.assertIn('Include MCP: True', md)
|
||||
tool_lines = [line for line in md.splitlines() if line.startswith('- ')]
|
||||
self.assertEqual(len(tool_lines), 0)
|
||||
|
||||
def test_tool_lines_contain_name_and_source_hint(self) -> None:
|
||||
tool = PortingModule(
|
||||
name='TestTool',
|
||||
responsibility='testing',
|
||||
source_hint='test/path.ts',
|
||||
)
|
||||
pool = ToolPool(tools=(tool,), simple_mode=False, include_mcp=False)
|
||||
md = pool.as_markdown()
|
||||
self.assertIn('- TestTool — test/path.ts', md)
|
||||
|
||||
|
||||
class TestToolPoolFrozen(unittest.TestCase):
|
||||
def test_cannot_mutate_fields(self) -> None:
|
||||
pool = ToolPool(tools=(), simple_mode=False, include_mcp=True)
|
||||
with self.assertRaises(AttributeError):
|
||||
pool.simple_mode = True # type: ignore[misc]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user