Orchestrator/shared/AI-Orchestrator/tests/test_orchestrator.py

"""Tests for orchestrator dispatch logic."""

from datetime import datetime, timezone

import pytest

from ai_orchestrator.models import BlockerRef, Issue, OrchestratorState, RunningEntry, LiveSession
from ai_orchestrator.orchestrator import Orchestrator


def make_issue(
    id: str = "id-1",
    identifier: str = "ABC-1",
    title: str = "Test",
    state: str = "Todo",
    priority: int | None = 1,
    created_at: datetime | None = None,
    blocked_by: list | None = None,
) -> Issue:
    return Issue(
        id=id,
        identifier=identifier,
        title=title,
        state=state,
        priority=priority,
        created_at=created_at or datetime(2024, 1, 1, tzinfo=timezone.utc),
        blocked_by=blocked_by or [],
    )


class TestDispatchSorting:
    """Tests for issue dispatch sorting."""

    def test_priority_ascending(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("---\ntracker:\n  kind: linear\n---\nprompt")
        orch = Orchestrator(workflow)

        issues = [
            make_issue(id="3", identifier="C", priority=3),
            make_issue(id="1", identifier="A", priority=1),
            make_issue(id="2", identifier="B", priority=2),
        ]
        sorted_issues = orch._sort_for_dispatch(issues)
        assert [i.identifier for i in sorted_issues] == ["A", "B", "C"]

    def test_null_priority_sorts_last(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("---\ntracker:\n  kind: linear\n---\nprompt")
        orch = Orchestrator(workflow)

        issues = [
            make_issue(id="1", identifier="A", priority=None),
            make_issue(id="2", identifier="B", priority=2),
        ]
        sorted_issues = orch._sort_for_dispatch(issues)
        assert sorted_issues[0].identifier == "B"
        assert sorted_issues[1].identifier == "A"

    def test_oldest_first_tiebreaker(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("---\ntracker:\n  kind: linear\n---\nprompt")
        orch = Orchestrator(workflow)

        issues = [
            make_issue(
                id="2", identifier="B", priority=1,
                created_at=datetime(2024, 2, 1, tzinfo=timezone.utc),
            ),
            make_issue(
                id="1", identifier="A", priority=1,
                created_at=datetime(2024, 1, 1, tzinfo=timezone.utc),
            ),
        ]
        sorted_issues = orch._sort_for_dispatch(issues)
        assert sorted_issues[0].identifier == "A"


class TestDispatchEligibility:
    """Tests for should_dispatch logic."""

    def setup_method(self, method, tmp_path=None):
        """Set up a basic orchestrator for testing."""
        pass

    def _make_orchestrator(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("""---
tracker:
  kind: linear
  project_slug: test
  api_key: fake-key
  active_states: ["Todo", "In Progress"]
  terminal_states: ["Done", "Cancelled"]
agent:
  max_concurrent_agents: 5
---
prompt""")
        orch = Orchestrator(workflow)
        orch._reload_workflow()
        return orch

    def test_active_issue_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(state="Todo")
        assert orch._should_dispatch(issue) is True

    def test_terminal_issue_not_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(state="Done")
        assert orch._should_dispatch(issue) is False

    def test_already_running_not_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(id="id-1", state="Todo")
        orch._state.running["id-1"] = RunningEntry(
            issue_id="id-1", identifier="ABC-1", issue=issue
        )
        assert orch._should_dispatch(issue) is False

    def test_already_claimed_not_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(id="id-1", state="Todo")
        orch._state.claimed.add("id-1")
        assert orch._should_dispatch(issue) is False

    def test_todo_with_non_terminal_blocker_not_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(
            state="Todo",
            blocked_by=[BlockerRef(id="b1", state="In Progress")],
        )
        assert orch._should_dispatch(issue) is False

    def test_todo_with_terminal_blocker_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(
            state="Todo",
            blocked_by=[BlockerRef(id="b1", state="Done")],
        )
        assert orch._should_dispatch(issue) is True

    def test_missing_required_fields_not_eligible(self, tmp_path):
        orch = self._make_orchestrator(tmp_path)
        issue = make_issue(title="", state="Todo")
        assert orch._should_dispatch(issue) is False


class TestBackoffComputation:
    """Tests for retry backoff calculation."""

    def test_first_retry(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("""---
tracker:
  kind: linear
  project_slug: test
  api_key: key
agent:
  max_retry_backoff_ms: 300000
---
prompt""")
        orch = Orchestrator(workflow)
        orch._reload_workflow()
        # attempt 1: 10000 * 2^0 = 10000
        assert orch._compute_backoff(1) == 10000

    def test_second_retry(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("""---
tracker:
  kind: linear
  project_slug: test
  api_key: key
agent:
  max_retry_backoff_ms: 300000
---
prompt""")
        orch = Orchestrator(workflow)
        orch._reload_workflow()
        # attempt 2: 10000 * 2^1 = 20000
        assert orch._compute_backoff(2) == 20000

    def test_backoff_capped(self, tmp_path):
        workflow = tmp_path / "WORKFLOW.md"
        workflow.write_text("""---
tracker:
  kind: linear
  project_slug: test
  api_key: key
agent:
  max_retry_backoff_ms: 300000
---
prompt""")
        orch = Orchestrator(workflow)
        orch._reload_workflow()
        # attempt 10: 10000 * 2^9 = 5120000, capped at 300000
        assert orch._compute_backoff(10) == 300000