mirror of
https://github.com/khodges42/nightShift.git
synced 2026-06-14 10:08:37 +00:00
242 lines
9.5 KiB
Python
242 lines
9.5 KiB
Python
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from nightshift.agents import AgentExecutor, build_prompt_bundle, parse_review_output, strip_ansi_escape_sequences
|
|
from nightshift.agents import AgentInvocation, format_agent_invocation
|
|
from nightshift.artifacts import ArtifactStore
|
|
from nightshift.config import AgentConfig, StageConfig
|
|
from nightshift.tasks import parse_tasks
|
|
|
|
|
|
TASK_MD = """# Tasks
|
|
|
|
- [ ] TASK-001: Add fake agent coverage
|
|
|
|
Description:
|
|
Exercise fake command agents.
|
|
|
|
Acceptance Criteria:
|
|
- Prompt includes task details
|
|
- Agent output is stored
|
|
"""
|
|
|
|
|
|
class AgentExecutorTests(unittest.TestCase):
|
|
def test_build_prompt_bundle_includes_task_and_acceptance_criteria(self) -> None:
|
|
task = parse_tasks(TASK_MD)[0]
|
|
prompt = build_prompt_bundle(
|
|
system_prompt="System rules",
|
|
stage=StageConfig(id="plan", type="agent", agent="planner"),
|
|
task=task,
|
|
project_context="Project context",
|
|
previous_outputs={"prior": "Earlier output"},
|
|
retry_notes=["Retry note"],
|
|
)
|
|
|
|
self.assertIn("System rules", prompt)
|
|
self.assertIn("TASK-001", prompt)
|
|
self.assertIn("- Prompt includes task details", prompt)
|
|
self.assertIn("Earlier output", prompt)
|
|
self.assertIn("Retry note", prompt)
|
|
|
|
def test_build_prompt_bundle_includes_task_context(self) -> None:
|
|
task = parse_tasks(TASK_MD)[0]
|
|
prompt = build_prompt_bundle(
|
|
system_prompt="System rules",
|
|
stage=StageConfig(id="plan", type="agent", agent="planner"),
|
|
task=task,
|
|
project_context="Project context",
|
|
task_context="Task context body",
|
|
previous_outputs={},
|
|
retry_notes=[],
|
|
retry_context="- No retries",
|
|
)
|
|
|
|
self.assertIn("## Task Context", prompt)
|
|
self.assertIn("Task context body", prompt)
|
|
self.assertIn("- No retries", prompt)
|
|
|
|
def test_file_writer_contract_mentions_repair_context(self) -> None:
|
|
task = parse_tasks(TASK_MD)[0]
|
|
prompt = build_prompt_bundle(
|
|
system_prompt="System rules",
|
|
stage=StageConfig(id="write", type="file_writer", agent="writer"),
|
|
task=task,
|
|
project_context="Project context",
|
|
previous_outputs={},
|
|
retry_notes=["Retry note"],
|
|
)
|
|
|
|
self.assertIn("On repair attempts", prompt)
|
|
self.assertIn("failed stage output", prompt)
|
|
self.assertIn("Use real project-relative paths", prompt)
|
|
self.assertNotIn("relative/path.py", prompt)
|
|
self.assertNotIn("including tests", prompt)
|
|
|
|
def test_file_writer_contract_includes_stage_allowed_paths(self) -> None:
|
|
task = parse_tasks(TASK_MD)[0]
|
|
prompt = build_prompt_bundle(
|
|
system_prompt="System rules",
|
|
stage=StageConfig(
|
|
id="write",
|
|
type="file_writer",
|
|
agent="writer",
|
|
allowed_paths=("story/chapters",),
|
|
),
|
|
project_context="Project context",
|
|
task=task,
|
|
previous_outputs={},
|
|
retry_notes=[],
|
|
)
|
|
|
|
self.assertIn("Use only paths under these project-relative targets: `story/chapters`.", prompt)
|
|
|
|
def test_command_agent_writes_output_and_returns_pass(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
prompt_path = root / "planner.md"
|
|
prompt_path.write_text("Plan carefully.", encoding="utf-8")
|
|
artifacts = ArtifactStore(root, ".nightshift", run_id="test-run")
|
|
executor = AgentExecutor(
|
|
root,
|
|
{
|
|
"planner": AgentConfig(
|
|
id="planner",
|
|
backend="command",
|
|
command='python -c "import sys; print(sys.stdin.read())"',
|
|
system_prompt=Path("planner.md"),
|
|
)
|
|
},
|
|
artifacts,
|
|
)
|
|
task = parse_tasks(TASK_MD)[0]
|
|
stage = StageConfig(id="plan", type="agent", agent="planner", output="plan.md")
|
|
|
|
result = executor.run_stage(stage, task)
|
|
|
|
self.assertEqual(result.status, "pass")
|
|
output = (root / result.output_path).read_text(encoding="utf-8")
|
|
self.assertIn("TASK-001", output)
|
|
self.assertIn("Plan carefully.", output)
|
|
|
|
def test_review_output_parser_accepts_structured_status(self) -> None:
|
|
status, reason, next_stage, context_update = parse_review_output(
|
|
"status: retry\nreason: Needs changes\nnext_stage: implement\ncontext_update: Fix tests\n"
|
|
)
|
|
|
|
self.assertEqual(status, "retry")
|
|
self.assertEqual(reason, "Needs changes")
|
|
self.assertEqual(next_stage, "implement")
|
|
self.assertEqual(context_update, "Fix tests")
|
|
|
|
def test_review_output_parser_treats_empty_sentinel_next_stage_as_missing(self) -> None:
|
|
for next_stage_value in ("", "None", "null", "N/A"):
|
|
with self.subTest(next_stage=next_stage_value):
|
|
status, reason, next_stage, context_update = parse_review_output(
|
|
f"status: pass\nreason: ok\nnext_stage: {next_stage_value}\ncontext_update: None\n"
|
|
)
|
|
|
|
self.assertEqual(status, "pass")
|
|
self.assertEqual(reason, "ok")
|
|
self.assertIsNone(next_stage)
|
|
self.assertIsNone(context_update)
|
|
|
|
def test_ollama_agent_invocation_uses_model_without_real_ollama(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
prompt_path = root / "planner.md"
|
|
prompt_path.write_text("Plan carefully.", encoding="utf-8")
|
|
artifacts = ArtifactStore(root, ".nightshift", run_id="test-run")
|
|
executor = AgentExecutor(
|
|
root,
|
|
{
|
|
"planner": AgentConfig(
|
|
id="planner",
|
|
backend="ollama",
|
|
command=None,
|
|
model="tiny-model",
|
|
system_prompt=Path("planner.md"),
|
|
)
|
|
},
|
|
artifacts,
|
|
)
|
|
task = parse_tasks(TASK_MD)[0]
|
|
stage = StageConfig(id="plan", type="agent", agent="planner", output="plan.md")
|
|
|
|
response = MagicMock()
|
|
response.__enter__.return_value.read.return_value = b'{"response":"ollama output"}'
|
|
|
|
with patch("nightshift.agents.request.urlopen", return_value=response) as urlopen:
|
|
result = executor.run_stage(stage, task)
|
|
|
|
self.assertEqual(result.status, "pass")
|
|
request_obj = urlopen.call_args.args[0]
|
|
body = request_obj.data.decode("utf-8")
|
|
self.assertIn('"model": "tiny-model"', body)
|
|
self.assertIn('"stream": false', body)
|
|
output = (root / result.output_path).read_text(encoding="utf-8")
|
|
self.assertIn("POST http://localhost:11434/api/generate", output)
|
|
self.assertIn("ollama output", output)
|
|
|
|
def test_openai_compatible_agent_sends_temperature(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
prompt_path = root / "planner.md"
|
|
prompt_path.write_text("Plan carefully.", encoding="utf-8")
|
|
artifacts = ArtifactStore(root, ".nightshift", run_id="test-run")
|
|
executor = AgentExecutor(
|
|
root,
|
|
{
|
|
"planner": AgentConfig(
|
|
id="planner",
|
|
backend="openai_compatible",
|
|
command=None,
|
|
model="tiny-model",
|
|
base_url="http://localhost:11434/v1",
|
|
temperature=0.2,
|
|
system_prompt=Path("planner.md"),
|
|
)
|
|
},
|
|
artifacts,
|
|
)
|
|
task = parse_tasks(TASK_MD)[0]
|
|
stage = StageConfig(id="plan", type="agent", agent="planner", output="plan.md")
|
|
response = MagicMock()
|
|
response.__enter__.return_value.read.return_value = (
|
|
b'{"choices":[{"message":{"content":"api output"}}]}'
|
|
)
|
|
|
|
with patch("nightshift.agents.request.urlopen", return_value=response) as urlopen:
|
|
result = executor.run_stage(stage, task)
|
|
|
|
self.assertEqual(result.status, "pass")
|
|
request_obj = urlopen.call_args.args[0]
|
|
body = request_obj.data.decode("utf-8")
|
|
self.assertIn('"temperature": 0.2', body)
|
|
self.assertIn("api output", (root / result.output_path).read_text(encoding="utf-8"))
|
|
|
|
def test_agent_artifact_format_tolerates_missing_streams(self) -> None:
|
|
invocation = AgentInvocation(
|
|
agent_id="planner",
|
|
command="ollama run model",
|
|
prompt="prompt",
|
|
exit_code=0,
|
|
stdout=None, # type: ignore[arg-type]
|
|
stderr=None, # type: ignore[arg-type]
|
|
duration_seconds=0.1,
|
|
)
|
|
|
|
output = format_agent_invocation("plan", invocation)
|
|
|
|
self.assertIn("Agent: `planner`", output)
|
|
self.assertIn("## stderr", output)
|
|
|
|
def test_strip_ansi_escape_sequences(self) -> None:
|
|
self.assertEqual(strip_ansi_escape_sequences("\x1b[?25lthinking\x1b[0m"), "thinking")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|