ChatAIBot/tests/test_rpg_story.py

import asyncio
import json
from unittest.mock import AsyncMock, patch

from services.rpg_story import (
    migrate_beats_to_steps,
    normalize_story_arc,
    get_current_step,
    step_progress,
    is_arc_completed,
    sync_quest_to_current_step,
    apply_step_advance,
    apply_story_post,
    is_new_arc_request,
    normalize_new_arc_first,
    new_arc_roll_choice,
    append_new_arc_roll_choice,
    format_step_guidance_for_character,
    format_new_arc_opening,
    normalize_step,
    resolve_step_injection,
    roll_next_arc,
)


def test_migrate_beats_to_steps():
    arc = {
        "beats": [
            {"id": "b1", "title": "Meet Yuki", "injection": "Door opens.", "choices": []},
        ]
    }
    out = migrate_beats_to_steps(arc)
    assert len(out["steps"]) == 1
    assert "beats" not in out
    assert out["steps"][0]["title"] == "Meet Yuki"


def test_step_progress():
    arc = normalize_story_arc({
        "steps": [{"id": "s1", "title": "A"}, {"id": "s2", "title": "B"}],
        "current_step_index": 1,
    })
    assert step_progress(arc) == (2, 2)


def test_is_arc_completed():
    arc = normalize_story_arc({
        "steps": [{"id": "s1", "title": "A"}],
        "current_step_index": 1,
        "status": "active",
    })
    assert is_arc_completed(arc)


def test_new_arc_request_detected():
    assert is_new_arc_request("Начать новую арку")
    assert is_new_arc_request("Начать новый квест: 'Тайна плазменных шаров'")
    assert is_new_arc_request("Start new quest")
    assert not is_new_arc_request("привет")


def test_normalize_new_arc_first():
    assert normalize_new_arc_first("user") == "user"
    assert normalize_new_arc_first("Character") == "character"
    assert normalize_new_arc_first("other") is None


def test_new_arc_roll_choice_and_filter():
    base = new_arc_roll_choice()
    assert base["type"] == "new_arc_roll"
    merged = append_new_arc_roll_choice([
        {"id": "x", "label": "Начать новый квест: Foo", "source": "narrator"},
        {"id": "y", "label": "Пойти спать", "source": "narrator"},
    ])
    assert len(merged) == 2
    assert merged[-1]["type"] == "new_arc_roll"
    assert not any("новый квест" in (c.get("label") or "").lower() for c in merged[:-1])


def test_normalize_step_repairs_legacy_beat_shape():
    step = normalize_step({
        "status_quo": "Сара и Григо прибывают к лаборатории.",
        "choices": [{"label": "Войти", "injection": "*Дверь скрипит*"}],
    }, 0)
    assert step["title"]
    assert step["completion_criteria"]
    assert resolve_step_injection(step) == "*Дверь скрипит*"


def test_format_new_arc_opening():
    arc = {"title": "Тайна", "global_story": "Расследование шаров."}
    step = {"status_quo": "У ворот лаборатории тихо.", "choices": []}
    text = format_new_arc_opening(arc, step, lang="ru")
    assert "Тайна" in text
    assert "лаборатории" in text or "Расследование" in text


def test_guidance_block_ru():
    step = {"title": "Приютить", "goal": "Впустить внутрь", "character_guidance": "Робость"}
    block = format_step_guidance_for_character(step, {"global_story": "x", "steps": [step]}, lang="ru")
    assert "Приютить" in block
    assert "ОБЯЗАТЕЛЬНОЕ" in block


def test_apply_step_advance():
    arc = normalize_story_arc({
        "title": "Test",
        "steps": [
            {"id": "s1", "title": "Step One", "choices": []},
            {"id": "s2", "title": "Step Two", "injection": "Next scene.", "choices": []},
        ],
        "current_step_index": 0,
    })

    async def run():
        with patch("services.memory.upsert_quest", new_callable=AsyncMock) as uq:
            with patch("services.memory.update_session_plot_arc", new_callable=AsyncMock):
                with patch(
                    "services.rpg_story.sync_quest_to_current_step",
                    new_callable=AsyncMock,
                ):
                    return await apply_step_advance("sess", arc), uq

    result, mock_uq = asyncio.run(run())
    assert result["advanced"] is True
    assert result["injection"] == "Next scene."
    mock_uq.assert_called()


def test_apply_story_post_step_complete_fallback_via_note():
    arc = normalize_story_arc({
        "title": "Test",
        "steps": [
            {"id": "s1", "title": "Step One", "choices": []},
            {"id": "s2", "title": "Step Two", "choices": []},
        ],
        "current_step_index": 0,
    })

    post = {
        "step_complete": False,
        "step_completion_note": "The scene clearly satisfies the completion criteria.",
    }

    async def run():
        with patch("services.memory.get_session", new_callable=AsyncMock) as gs:
            gs.return_value = {"plot_arc_json": json.dumps(arc, ensure_ascii=False)}
            with patch(
                "services.rpg_story.apply_step_advance",
                new_callable=AsyncMock,
            ) as asa:
                asa.return_value = {"advanced": True, "arc_completed": False, "new_step": arc["steps"][1], "injection": ""}
                r = await apply_story_post("sess", post, arc, {"quests": True})
                return r, asa

    r, asa = asyncio.run(run())
    assert r["step_advanced"] is True
    assert asa.await_count == 1


def test_roll_next_arc_combines_context_for_generate_next_arc():
    old_arc = normalize_story_arc({
        "title": "Arc One",
        "steps": [{"id": "s1", "title": "Done"}],
        "current_step_index": 1,
        "status": "completed",
    })
    new_arc = normalize_story_arc({
        "title": "Arc Two",
        "steps": [{"id": "s1", "title": "Start"}],
        "current_step_index": 0,
    })

    async def run():
        with patch("services.memory.get_session", new_callable=AsyncMock) as gs:
            gs.return_value = {"plot_arc_json": json.dumps(old_arc, ensure_ascii=False)}
            with patch("services.rpg_plot.generate_next_arc", new_callable=AsyncMock) as gna:
                gna.return_value = new_arc
                with patch("services.memory.update_session_plot_arc", new_callable=AsyncMock):
                    with patch(
                        "services.rpg_story.sync_quest_to_current_step",
                        new_callable=AsyncMock,
                    ):
                        return await roll_next_arc(
                            "sess",
                            {"name": "Sarah", "description": "", "scenario": ""},
                            "Начать новую арку",
                            "adventure",
                            recent_context="assistant: goodbye",
                        ), gna

    result, gna = asyncio.run(run())
    assert result["title"] == "Arc Two"
    gna.assert_awaited_once()
    assert gna.await_args.args[3] == "assistant: goodbye\nНачать новую арку"