ChatAIBot/tests/test_arc_stuck_recovery.py

import asyncio
from unittest.mock import AsyncMock, patch

from services.rpg_plot import (
    prune_beats_for_done_quests,
    process_arc_beats,
    should_advance_arc_keywords,
    pop_matching_beats,
    dice_outcome_to_beat_trigger,
    active_quest_titles_to_close,
    can_fire_beat,
    record_beat_fired,
    complete_quest_for_fired_beat,
)


def test_prune_removes_beat_when_quest_done():
    arc = {
        "beats": [
            {"id": "b2", "title": "Highway Howl", "trigger": "event_driven:travel", "choices": []},
        ]
    }
    quests = [{"title": "Highway Howl", "status": "done"}]
    arc, removed = prune_beats_for_done_quests(arc, quests)
    assert removed[0]["title"] == "Highway Howl"
    assert arc["beats"] == []


def test_stuck_recovery_fires_when_no_active_quests_and_cooldown_ok():
    arc = {
        "beats": [
            {"id": "b3", "title": "New Beat", "trigger": "event_driven:rest", "choices": [{"id": "a", "label": "A"}]},
        ],
        "meta": {"last_beat_fired_at_user_turn": 1},
    }
    quests = [{"title": "Old", "status": "done"}]

    async def run():
        return await process_arc_beats(arc, quests, "hello", user_turn=5)

    arc2, fired, pruned, mode, extras = asyncio.run(run())
    assert mode == "stuck_recovery"
    assert fired[0]["title"] == "New Beat"
    assert arc2["beats"] == []
    assert extras.get("cooldown_skipped") is False


def test_stuck_recovery_skipped_when_cooldown_active():
    arc = {
        "beats": [
            {"id": "b3", "title": "New Beat", "trigger": "event_driven:rest", "choices": []},
        ],
        "meta": {"last_beat_fired_at_user_turn": 4},
    }
    quests = []

    async def run():
        return await process_arc_beats(arc, quests, "hello", user_turn=5)

    arc2, fired, _, mode, extras = asyncio.run(run())
    assert mode == ""
    assert fired == []
    assert extras.get("cooldown_skipped") is True
    assert len(arc2["beats"]) == 1


def test_stuck_recovery_skipped_after_reconcile_closed_quests():
    arc = {
        "beats": [
            {"id": "b3", "title": "New Beat", "trigger": "event_driven:rest", "choices": []},
        ],
    }
    quests = []

    async def run():
        return await process_arc_beats(
            arc, quests, "hello", user_turn=10, reconcile_closed_count=2
        )

    _, fired, _, mode, _ = asyncio.run(run())
    assert mode == ""
    assert fired == []


def test_dice_outcome_maps_to_after_fail():
    assert dice_outcome_to_beat_trigger("failure") == "event_driven:after_fail"
    assert dice_outcome_to_beat_trigger("critical failure") == "event_driven:after_fail"
    assert dice_outcome_to_beat_trigger("success") == "event_driven:after_success"


def test_after_fail_beat_fires_on_dice_failure_only_with_needs_check():
    arc = {
        "beats": [
            {
                "id": "b_fail",
                "title": "Dust Yourself Off",
                "trigger": "event_driven:after_fail",
                "injection": "The stumble leaves you both shaken.",
                "choices": [{"id": "a", "label": "Try again"}],
            },
        ]
    }

    async def run(needs_check: bool):
        return await process_arc_beats(
            arc, [], "продолжаем", last_dice_outcome="failure", needs_check=needs_check, user_turn=3
        )

    arc_ok, fired, _, mode, _ = asyncio.run(run(True))
    assert mode == "after_dice"
    assert fired[0]["id"] == "b_fail"
    assert arc_ok["beats"] == []

    arc2 = {
        "beats": [
            {
                "id": "b_fail",
                "title": "Dust Yourself Off",
                "trigger": "event_driven:after_fail",
                "choices": [],
            },
        ]
    }

    async def run_no_check():
        with patch(
            "services.rpg_plot.classify_plot_beat",
            new_callable=AsyncMock,
            return_value=None,
        ):
            return await process_arc_beats(
                arc2,
                [],
                "продолжаем",
                last_dice_outcome="failure",
                needs_check=False,
                user_turn=3,
                allow_stuck_recovery=False,
            )

    _, fired2, _, mode2, _ = asyncio.run(run_no_check())
    assert mode2 == ""
    assert fired2 == []


def test_orphan_active_quests_closed_when_not_in_arc():
    arc = {
        "beats": [
            {
                "id": "b_new_3",
                "title": "Defensive Magic",
                "trigger": "event_driven:after_fail",
                "choices": [],
            }
        ]
    }
    quests = [
        {"title": "Floral Whispers", "status": "active"},
        {"title": "Gift from the Glade", "status": "active"},
        {"title": "Defensive Magic", "status": "active"},
        {"title": "Old Done", "status": "done"},
    ]
    closed = active_quest_titles_to_close(arc, quests)
    assert set(closed) == {"Floral Whispers", "Gift from the Glade"}


def test_can_fire_beat_respects_min_gap():
    arc = {"meta": {"last_beat_fired_at_user_turn": 5}}
    assert can_fire_beat(arc, 6, min_gap=2) is False
    assert can_fire_beat(arc, 7, min_gap=2) is True


def test_record_beat_fired_updates_meta():
    arc = {}
    record_beat_fired(arc, {"id": "b1"}, 8)
    assert arc["meta"]["last_beat_fired_at_user_turn"] == 8
    assert arc["meta"]["last_beat_id"] == "b1"


def test_complete_quest_for_fired_beat():
    beat = {"id": "b1", "title": "Pancake Talk", "injection": "..."}

    async def run():
        with patch("services.memory.upsert_quest", new_callable=AsyncMock) as mock:
            await complete_quest_for_fired_beat("sess-1", beat)
            mock.assert_called_once_with("sess-1", "Pancake Talk", "done")

    asyncio.run(run())


def test_keyword_fallback_travel():
    arc = {
        "beats": [
            {"id": "b2", "title": "Highway Howl", "trigger": "event_driven:travel", "choices": []},
        ]
    }
    trig = should_advance_arc_keywords("едем на стадион")
    assert trig == "event_driven:travel"
    arc, fired = pop_matching_beats(arc, trig, max_beats=1)
    assert fired[0]["title"] == "Highway Howl"