Agent Test Harnesses

class MockToolExecutor:
    def __init__(self):
        self.mock_responses = {}
        self.call_log = []

    def register_mock(self, tool_name, arguments_pattern, response):
        """Register a mock response for a tool call."""
        self.mock_responses[(tool_name, arguments_pattern)] = response

    def execute(self, tool_name, arguments):
        """Return mock response instead of calling real tool."""
        self.call_log.append({
            "tool": tool_name,
            "arguments": arguments,
            "timestamp": time.time()
        })

        # Find matching mock
        for (name, pattern), response in self.mock_responses.items():
            if name == tool_name and self.matches(pattern, arguments):
                return response

        return {"error": f"No mock registered for {tool_name}"}

    def assert_tool_called(self, tool_name, expected_args=None):
        """Assert that a tool was called during the test."""
        calls = [c for c in self.call_log if c["tool"] == tool_name]
        assert len(calls) > 0, f"Expected {tool_name} to be called"
        if expected_args:
            assert any(self.matches(expected_args, c["arguments"]) for c in calls)

# Usage in tests
mock = MockToolExecutor()
mock.register_mock("kubectl_get", {"resource": "pods"}, {"pods": [...]})
mock.register_mock("query_prometheus", {"query": "rate(.*)"}, {"values": [...]})

agent = OpsAgent(tools=mock)
result = agent.run("Why is checkout slow?")

mock.assert_tool_called("query_prometheus")
mock.assert_tool_called("kubectl_get", {"resource": "pods"})

golden_transcript = {
    "id": "incident-redis-failover-001",
    "input": {
        "alert": "Redis connection timeout on checkout service",
        "context": "Production environment, Redis cluster with 3 replicas"
    },
    "expected_tool_sequence": [
        {"tool": "query_prometheus", "args": {"query": "redis_connected_clients"}},
        {"tool": "kubectl_get", "args": {"resource": "pods", "label_selector": "app=redis"}},
        {"tool": "query_prometheus", "args": {"query": "redis_up"}},
    ],
    "expected_final_answer_contains": ["Redis", "failover", "replica"],
    "max_tool_calls": 5,
    "max_tokens": 3000,
}

def verify_transcript(agent, transcript):
    """Run the agent and compare against golden transcript."""
    mock = MockToolExecutor()
    # Register mocks for expected tool calls
    agent.tools = mock

    result = agent.run(transcript["input"]["alert"])

    # Verify tool sequence
    actual_sequence = [c["tool"] for c in mock.call_log]
    expected_sequence = [t["tool"] for t in transcript["expected_tool_sequence"]]

    assert actual_sequence == expected_sequence, (
        f"Tool sequence mismatch:\nExpected: {expected_sequence}\nActual: {actual_sequence}"
    )

    # Verify final answer
    for term in transcript["expected_final_answer_contains"]:
        assert term.lower() in result.lower(), f"Answer missing: {term}"

class FailureInjector:
    def __init__(self):
        self.failures = {}

    def inject(self, tool_name, failure_type, probability=1.0):
        """Inject a failure for a specific tool."""
        self.failures[tool_name] = {
            "type": failure_type,
            "probability": probability
        }

    def execute_with_failures(self, tool_name, arguments):
        """Execute a tool with injected failures."""
        if tool_name in self.failures:
            failure = self.failures[tool_name]
            if random.random() < failure["probability"]:
                if failure["type"] == "timeout":
                    raise TimeoutError(f"{tool_name} timed out")
                elif failure["type"] == "permission":
                    raise PermissionError(f"Access denied to {tool_name}")
                elif failure["type"] == "error":
                    raise Exception(f"{tool_name} returned an error")

        return real_execute(tool_name, arguments)

# Test agent handles failures gracefully
injector = FailureInjector()
injector.inject("query_prometheus", "timeout", probability=0.5)

# Agent should retry or fall back
result = agent.run("Check Redis status")
assert result is not None, "Agent should handle tool failures"

import pytest

@pytest.mark.parametrize("transcript", golden_transcripts, ids=lambda t: t["id"])
def test_agent_behavior(transcript):
    """Test agent against golden transcripts."""
    agent = create_test_agent()
    verify_transcript(agent, transcript)

@pytest.mark.parametrize("failure_case", failure_test_cases)
def test_agent_resilience(failure_case):
    """Test agent handles failures gracefully."""
    agent = create_test_agent()
    injector = FailureInjector()

    for tool, failure in failure_case["failures"].items():
        injector.inject(tool, failure["type"], failure["probability"])

    result = agent.run(failure_case["input"])
    assert result is not None
    assert failure_case["expected_behavior"] in result

def test_agent_cost_budget(transcript):
    """Test agent stays within token budget."""
    agent = create_test_agent()
    result = agent.run(transcript["input"])

    assert result["total_tokens"] <= transcript["max_tokens"],
        f"Exceeded token budget: {result['total_tokens']} > {transcript['max_tokens']}"