"""Tests for orchestration.py token tracking integration (issue #634). Verifies: - log_token_usage writes to JSONL - log_token_usage calls token_budget.record_usage - check_budget enforces limits - Huey signal hook fires on task completion - Pipeline tasks are registered """ import json import os import tempfile from pathlib import Path from unittest.mock import MagicMock, patch import pytest class TestLogTokenUsage: """Test log_token_usage function.""" def test_skips_non_dict_result(self): """Should silently skip non-dict results.""" from orchestration import log_token_usage # Should not raise log_token_usage("test_task", None) log_token_usage("test_task", "string") log_token_usage("test_task", 42) def test_skips_zero_tokens(self): """Should skip entries with zero tokens.""" from orchestration import log_token_usage with patch("orchestration.TOKEN_LOG") as mock_log: mock_log.parent = MagicMock() log_token_usage("test_task", {"input_tokens": 0, "output_tokens": 0}) # Should not write to file mock_log.parent.mkdir.assert_not_called() def test_writes_to_jsonl(self, tmp_path): """Should append token usage to JSONL log.""" log_file = tmp_path / "token_usage.jsonl" with patch("orchestration.TOKEN_LOG", log_file), patch("orchestration.record_usage"): from orchestration import log_token_usage log_token_usage("playground_factory_task", { "input_tokens": 100, "output_tokens": 200, }) assert log_file.exists() line = json.loads(log_file.read_text().strip()) assert line["pipeline"] == "playground-factory" assert line["input_tokens"] == 100 assert line["output_tokens"] == 200 assert line["total_tokens"] == 300 def test_calls_budget_record_usage(self, tmp_path): """Should call token_budget.record_usage for budget tracking.""" log_file = tmp_path / "token_usage.jsonl" mock_record = MagicMock(return_value={"daily_remaining": 400000}) with patch("orchestration.TOKEN_LOG", log_file), patch("orchestration.record_usage", mock_record): from orchestration import log_token_usage log_token_usage("training_factory_task", { "input_tokens": 500, "output_tokens": 1000, }) mock_record.assert_called_once_with("training-factory", 500, 1000) def test_pipeline_name_derived_from_task(self, tmp_path): """Pipeline name should strip _task suffix and use hyphens.""" log_file = tmp_path / "token_usage.jsonl" with patch("orchestration.TOKEN_LOG", log_file), patch("orchestration.record_usage"): from orchestration import log_token_usage log_token_usage("knowledge_mine_task", { "input_tokens": 10, "output_tokens": 20, }) line = json.loads(log_file.read_text().strip()) assert line["pipeline"] == "knowledge-mine" class TestCheckBudget: """Test check_budget function.""" def test_returns_true_when_budget_available(self): """Should return True when can_afford returns True.""" with patch("orchestration.can_afford", return_value=True): from orchestration import check_budget assert check_budget("test", 1000) is True def test_returns_false_when_budget_exhausted(self): """Should return False when can_afford returns False.""" with patch("orchestration.can_afford", return_value=False), patch("orchestration.get_remaining", return_value=50): from orchestration import check_budget assert check_budget("test", 10000) is False def test_returns_true_when_budget_module_missing(self): """Should return True (no enforcement) when token_budget not importable.""" with patch("orchestration.can_afford", side_effect=ImportError): from orchestration import check_budget assert check_budget("test", 999999) is True class TestPipelineTasks: """Test Huey pipeline task registration.""" def test_all_pipelines_registered(self): """All 5 pipeline tasks should be registered with Huey.""" from orchestration import ( playground_factory_task, training_factory_task, knowledge_mine_task, adversary_task, codebase_genome_task, ) tasks = [ playground_factory_task, training_factory_task, knowledge_mine_task, adversary_task, codebase_genome_task, ] for task in tasks: # Huey tasks have a .call_local method assert hasattr(task, "call_local"), f"{task.__name__} not registered with Huey" def test_run_pipeline_returns_structured_result(self, tmp_path): """_run_pipeline should return dict with pipeline, status, tokens.""" # Create a stub script stub = tmp_path / "stub.sh" stub.write_text("#!/bin/bash\necho 'tokens used: 42'\n") stub.chmod(0o755) with patch("orchestration.check_budget", return_value=True): from orchestration import _run_pipeline result = _run_pipeline("test-pipeline", str(stub), 1000) assert result["pipeline"] == "test-pipeline" assert result["status"] == "success" assert "input_tokens" in result assert "output_tokens" in result def test_run_pipeline_skips_when_budget_exhausted(self): """Should return skipped status when budget is exhausted.""" with patch("orchestration.check_budget", return_value=False): from orchestration import _run_pipeline result = _run_pipeline("expensive", "/nonexistent", 999999) assert result["status"] == "skipped" assert result["reason"] == "budget_exhausted" def test_run_pipeline_handles_missing_script(self): """Should return failed status when script doesn't exist.""" with patch("orchestration.check_budget", return_value=True): from orchestration import _run_pipeline result = _run_pipeline("broken", "/no/such/script.sh", 1000) assert result["status"] == "failed" assert result["reason"] == "script_not_found"