orchestrator/tests/test_routing_golden.py

"""Golden-set routing test for the structured router.

20 labelled tasks → expected agent. Live tests require ANTHROPIC_API_KEY +
COMPOSIO_API_KEY; without them the parametrized cases skip at runtime but
the module still collects (so `pytest --collect-only` in CI works without
secrets). Run with:

    pytest tests/test_routing_golden.py -v
"""

import os

import pytest
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))

_REQUIRED_KEYS = ("ANTHROPIC_API_KEY", "COMPOSIO_API_KEY")
_missing = [k for k in _REQUIRED_KEYS if not os.getenv(k)]
requires_keys = pytest.mark.skipif(
    bool(_missing),
    reason=f"Requires {', '.join(_REQUIRED_KEYS)}; missing: {', '.join(_missing)}",
)


GOLDEN_SET: list[tuple[str, str]] = [
    # implementer (3)
    (
        "Write a Python function that validates email addresses using a regex",
        "implementer",
    ),
    (
        "Add a new Lambda handler in handlers/notify.py that publishes an SNS message",
        "implementer",
    ),
    (
        "Fix the bug in our auth middleware where expired tokens are accepted as valid",
        "implementer",
    ),
    # reviewer (3)
    (
        "Review this pull request diff for correctness, security, and maintainability concerns",
        "reviewer",
    ),
    (
        "Code review the attached commit and categorize each issue as BLOCK, FIX, or NIT",
        "reviewer",
    ),
    (
        "Review the following code changes and tell me what should block merge",
        "reviewer",
    ),
    # researcher (3)
    (
        "What is the latest stable version of the langgraph Python package?",
        "researcher",
    ),
    (
        "Look up the AWS Lambda maximum concurrent execution limit in us-east-1",
        "researcher",
    ),
    ("Find documentation on how to configure DynamoDB TTL", "researcher"),
    # cross_reviewer (2)
    (
        "Get a cross-family second opinion on this diff using GPT to catch what Claude might miss",
        "cross_reviewer",
    ),
    (
        "Run an independent cross-model review on this Lambda handler change",
        "cross_reviewer",
    ),
    # scanner (3)
    (
        "Scan the entire monorepo to find inconsistent error handling patterns across 200+ files",
        "scanner",
    ),
    (
        "Analyze the whole codebase for unused imports and dead code using a large-context model",
        "scanner",
    ),
    ("Audit every Lambda handler in this repo for hardcoded secrets", "scanner"),
    # fast_coder (3)
    (
        "Quick small task using DeepSeek: write a 5-line Python helper that converts kebab-case to snake_case",
        "fast_coder",
    ),
    (
        "Fast bounded coding job: implement a one-function utility that pads strings to a fixed width",
        "fast_coder",
    ),
    (
        "Use the cheap fast coder to write a short Python snippet that parses a CSV row into a dict",
        "fast_coder",
    ),
    # connector (3)
    ("Send a Slack message to the #ops channel announcing the deploy", "connector"),
    ("Create a Notion page under the Engineering space called 'Q3 plan'", "connector"),
    ("Find a file named contracts.pdf in my Google Drive", "connector"),
]


def test_golden_set_size():
    assert len(GOLDEN_SET) == 20


@requires_keys
@pytest.mark.parametrize("task,expected", GOLDEN_SET)
def test_router_picks_expected_agent(task: str, expected: str):
    # Lazy import — graph.py loads composio_tools at module import and needs
    # COMPOSIO_API_KEY. Keep the import inside the skipif-guarded test body.
    from graph import router_node

    out = router_node({"task": task, "messages": []})
    assert out["route"] == expected, (
        f"task={task!r} got={out['route']} expected={expected}"
    )