mirror of
https://github.com/khodges42/nightShift.git
synced 2026-06-14 18:18:36 +00:00
114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
from pathlib import Path
|
|
from dataclasses import replace
|
|
import tempfile
|
|
import unittest
|
|
|
|
from nightshift.artifacts import ArtifactStore
|
|
from nightshift.config import SafetyConfig, StageConfig
|
|
from nightshift.pipeline import PipelineRunner
|
|
from nightshift.semantic_index import build_semantic_index, search_index
|
|
from nightshift.tasks import parse_tasks
|
|
from nightshift.telemetry import estimate_tokens, format_telemetry_summary, telemetry_from_stage_output
|
|
|
|
from tests.test_pipeline import TASK_MD, make_config, _write_common_files
|
|
|
|
|
|
class TelemetryAndIndexTests(unittest.TestCase):
|
|
def test_telemetry_estimates_tokens_and_groups_by_model(self) -> None:
|
|
output = "\n".join(
|
|
[
|
|
"# Agent Output: plan",
|
|
"",
|
|
"Agent: `planner`",
|
|
"Duration seconds: 1.250",
|
|
"",
|
|
"## stdout",
|
|
"",
|
|
"```text",
|
|
"plan ok",
|
|
"```",
|
|
"",
|
|
"## Prompt",
|
|
"",
|
|
"```markdown",
|
|
"hello world",
|
|
"```",
|
|
]
|
|
)
|
|
|
|
entry = telemetry_from_stage_output(
|
|
stage_id="plan",
|
|
stage_type="agent",
|
|
status="pass",
|
|
output=output,
|
|
retry_count=0,
|
|
model="qwen2.5-coder:14b",
|
|
)
|
|
summary = format_telemetry_summary((entry,))
|
|
|
|
self.assertGreater(estimate_tokens("hello world"), 0)
|
|
self.assertEqual(entry.agent_id, "planner")
|
|
self.assertEqual(entry.duration_seconds, 1.25)
|
|
self.assertIn("qwen2.5-coder:14b", summary)
|
|
|
|
def test_pipeline_writes_telemetry_summary(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
_write_common_files(root)
|
|
stages = (StageConfig(id="plan", type="agent", agent="planner", output="plan.md"),)
|
|
config = make_config(root, stages)
|
|
runner = PipelineRunner(config, ArtifactStore(root, ".nightshift", run_id="test-run"))
|
|
|
|
result = runner.run_task(parse_tasks(TASK_MD)[0])
|
|
|
|
task_dir = root / ".nightshift" / "runs" / "test-run" / "tasks" / "TASK-001"
|
|
self.assertEqual(result.status, "complete")
|
|
self.assertTrue((task_dir / "telemetry-summary.md").exists())
|
|
self.assertTrue((root / ".nightshift" / "runs" / "test-run" / "telemetry-summary.md").exists())
|
|
|
|
def test_semantic_index_finds_symbols_and_tests(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
(root / "src").mkdir()
|
|
(root / "tests").mkdir()
|
|
(root / "src" / "service.py").write_text(
|
|
"import sqlite3\n\nclass SnippetStore:\n pass\n\ndef create_snippet():\n return True\n",
|
|
encoding="utf-8",
|
|
)
|
|
(root / "tests" / "test_service.py").write_text(
|
|
"def test_create_snippet():\n assert True\n",
|
|
encoding="utf-8",
|
|
)
|
|
safety = SafetyConfig(
|
|
require_clean_worktree=False,
|
|
scoped_paths=("src", "tests"),
|
|
allowed_commands=(),
|
|
forbidden_commands=(),
|
|
)
|
|
|
|
index = build_semantic_index(root, safety)
|
|
results = search_index(index, "create snippet sqlite")
|
|
|
|
self.assertTrue(any("create_snippet" in item.symbols for item in index))
|
|
self.assertTrue(any(item.path == "src/service.py" for item in results))
|
|
|
|
def test_semantic_context_stage_writes_artifacts(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
_write_common_files(root)
|
|
(root / "snippet.py").write_text("def create_snippet():\n return 'ok'\n", encoding="utf-8")
|
|
stages = (StageConfig(id="semantic", type="semantic_context", output="semantic-context.md"),)
|
|
config = make_config(root, stages)
|
|
runner = PipelineRunner(config, ArtifactStore(root, ".nightshift", run_id="test-run"))
|
|
|
|
result = runner.run_task(parse_tasks(TASK_MD)[0])
|
|
|
|
task_dir = root / ".nightshift" / "runs" / "test-run" / "tasks" / "TASK-001"
|
|
self.assertEqual(result.status, "complete")
|
|
self.assertTrue((task_dir / "semantic-index.md").exists())
|
|
self.assertTrue((task_dir / "semantic-context.md").exists())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|