tests/test_generator.py

import inspect
import json

from turnabout.generation.generator import CaseGenerator


class _Message:
    def __init__(self, content):
        self.content = content


class _Choice:
    def __init__(self, content):
        self.message = _Message(content)


class _Completion:
    def __init__(self, content):
        self.choices = [_Choice(content)]


class _Completions:
    def __init__(self, payloads):
        self._payloads = list(payloads)

    def create(self, **_kwargs):
        payload = self._payloads.pop(0)
        return _Completion("```json\n" + json.dumps(payload) + "\n```")


class _Chat:
    def __init__(self, payloads):
        self.completions = _Completions(payloads)


class _OpenAIClient:
    def __init__(self, payloads):
        self.chat = _Chat(payloads)


class APITimeoutError(Exception):
    pass


class _TimeoutCompletions:
    def __init__(self):
        self.calls = 0

    def create(self, **_kwargs):
        self.calls += 1
        raise APITimeoutError("request timed out")


class _TimeoutChat:
    def __init__(self):
        self.completions = _TimeoutCompletions()


class _TimeoutOpenAIClient:
    def __init__(self):
        self.chat = _TimeoutChat()


class RateLimitError(Exception):
    status_code = 429


class _RateLimitCompletions:
    def __init__(self):
        self.calls = 0

    def create(self, **_kwargs):
        self.calls += 1
        raise RateLimitError("Error code: 429 - Too many requests")


class _RateLimitChat:
    def __init__(self):
        self.completions = _RateLimitCompletions()


class _RateLimitOpenAIClient:
    def __init__(self):
        self.chat = _RateLimitChat()


def _easy_case_payloads():
    return [
        {
            "title": "The Clocktower Switch",
            "description": "A staged theft hides a simple timing contradiction.",
            "defendant": "d_defendant",
            "true_culprit": "w_witness",
            "witnesses": ["w_witness"],
            "crime_type": "theft",
        },
        {
            "characters": [
                {
                    "id": "d_defendant",
                    "name": "Dana Vale",
                    "description": "The accused apprentice.",
                    "role": "defendant",
                },
                {
                    "id": "w_witness",
                    "name": "Milo Kent",
                    "description": "The witness who lied about the time.",
                    "role": "witness",
                },
            ],
            "evidence": [
                {
                    "id": "clock_photo",
                    "name": "Clock Photo",
                    "item_type": "photo",
                    "description": "A photo of the tower clock.",
                    "detail": "The clock shows 8:10, not 8:30.",
                    "source": "pre_given",
                },
                {
                    "id": "repair_log",
                    "name": "Repair Log",
                    "item_type": "document",
                    "description": "A log of the tower maintenance.",
                    "detail": "The clock was stopped for repairs at 8:15.",
                    "source": "pre_given",
                },
            ],
        },
        {
            "court": {
                "penalty_limit": 5,
                "rounds": [
                    {
                        "witness_id": "w_witness",
                        "order": 1,
                        "initial_testimony_id": "testimony_1",
                        "testimonies": [
                            {
                                "id": "testimony_1",
                                "title": "What I Saw",
                                "witness_id": "w_witness",
                                "preamble": "The witness describes the scene.",
                                "statements": [
                                    {
                                        "id": "stmt_time",
                                        "text": "I saw Dana at the tower at 8:30.",
                                        "press_response": "I am certain about the time.",
                                        "contradiction": {
                                            "evidence_id": "clock_photo",
                                            "explanation": (
                                                "The clock photo proves the "
                                                "witness saw a different time."
                                            ),
                                            "is_primary": True,
                                        },
                                    }
                                ],
                            }
                        ],
                    }
                ],
                "win_condition": {
                    "required_contradiction_ids": ["stmt_time"],
                    "final_evidence_id": "clock_photo",
                },
            }
        },
    ]


def _hard_case_payloads_with_investigation_retry():
    payloads = _easy_case_payloads()
    payloads[1]["evidence"].append(
        {
            "id": "kitchen_access_log",
            "name": "Kitchen Access Log",
            "item_type": "document",
            "description": "A record of kitchen door access.",
            "detail": "The witness entered the kitchen before the theft.",
            "source": "investigation",
            "acquisition": {
                "method": "examine",
                "location_id": "main_kitchen",
                "target_id": "kitchen_access_log_terminal",
                "prerequisites": [],
                "present_evidence_id": None,
            },
        }
    )
    payloads.append(
        {
            "locations": [
                {
                    "id": "front_counter",
                    "name": "Front Counter",
                    "description": "The public counter area.",
                    "connected_to": [],
                    "examinables": [],
                    "characters_present": [],
                    "available_from_start": True,
                    "unlock_prerequisites": [],
                }
            ],
            "dialogues": [],
            "investigation_gate": {
                "required_evidence": ["kitchen_access_log"],
                "auto_advance": False,
            },
        }
    )
    payloads.append(
        {
            "locations": [
                {
                    "id": "main_kitchen",
                    "name": "Main Kitchen",
                    "description": "The kitchen where the access terminal sits.",
                    "connected_to": [],
                    "examinables": [
                        {
                            "id": "kitchen_access_log_terminal",
                            "name": "Kitchen Access Log Terminal",
                            "description": "A terminal storing access records.",
                            "examined_description": "It shows the witness entered first.",
                        }
                    ],
                    "characters_present": [],
                    "available_from_start": True,
                    "unlock_prerequisites": [],
                }
            ],
            "dialogues": [],
            "investigation_gate": {
                "required_evidence": ["kitchen_access_log"],
                "auto_advance": False,
            },
        }
    )
    return payloads


def test_generate_is_sync_and_reports_progress():
    generator = CaseGenerator(
        llm_client=_OpenAIClient(_easy_case_payloads()),
        backend="openai",
    )
    progress_updates = []

    assert not inspect.iscoroutinefunction(generator.generate)
    assert not inspect.iscoroutinefunction(generator._step_with_retry)

    case = generator.generate(
        theme="clocktower",
        difficulty="easy",
        progress_callback=lambda value, desc: progress_updates.append(
            (value, desc)
        ),
    )

    assert case.title == "The Clocktower Switch"
    assert case.court.win_condition.required_contradiction_ids == ["stmt_time"]
    assert progress_updates[0] == (0.25, "Generating narrative...")
    assert progress_updates[-2] == (0.9, "Assembling and validating case...")
    assert progress_updates[-1] == (1.0, "Case generated successfully.")


def test_generate_events_yields_status_before_final_case():
    generator = CaseGenerator(
        llm_client=_OpenAIClient(_easy_case_payloads()),
        backend="openai",
    )

    events = list(generator.generate_events(theme="clocktower", difficulty="easy"))

    assert events[0] == (0.25, "Generating narrative...", None)
    assert events[-1][0] == 1.0
    assert events[-1][1] == "Case generated successfully."
    assert events[-1][2].title == "The Clocktower Switch"


def test_generate_timeout_is_not_retried():
    client = _TimeoutOpenAIClient()
    generator = CaseGenerator(
        llm_client=client,
        backend="openai",
        request_timeout=1,
    )

    try:
        generator.generate(theme="clocktower", difficulty="easy")
    except TimeoutError as exc:
        assert "LLM request timed out during step 'narrative'" in str(exc)
    else:
        raise AssertionError("generate() should raise TimeoutError")

    assert client.chat.completions.calls == 1


def test_generate_rate_limit_is_not_retried():
    client = _RateLimitOpenAIClient()
    generator = CaseGenerator(
        llm_client=client,
        backend="openai",
        request_timeout=1,
    )

    try:
        generator.generate(theme="clocktower", difficulty="easy")
    except RuntimeError as exc:
        assert "LLM gateway rate limit reached during step 'narrative'" in str(exc)
        assert "LLM_STEP_DELAY_SECONDS" in str(exc)
    else:
        raise AssertionError("generate() should raise RuntimeError")

    assert client.chat.completions.calls == 1


def test_hard_generation_retries_invalid_investigation_without_repair():
    generator = CaseGenerator(
        llm_client=_OpenAIClient(_hard_case_payloads_with_investigation_retry()),
        backend="openai",
    )

    case = generator.generate(theme="bakery", difficulty="hard")

    main_kitchen = case.get_location("main_kitchen")
    assert main_kitchen is not None
    assert case.get_location("front_counter") is None
    assert any(
        ex.id == "kitchen_access_log_terminal"
        for ex in main_kitchen.examinables
    )