WalkXR-AI / tests / test_memory_simulation.py
test_memory_simulation.py
Raw
"""
Unit tests for memory-focused simulation logic (T3.6 Part 1 / 2).

Covers:
- Memory scenario YAML loads correctly into SimulationCase objects
- scripted_user_input() injects the recall question on the configured turn
- check_memory_recall() returns hit/miss correctly for single- and multi-detail
  recall scenarios
- check_continuity() flags obvious continuity-break phrasing

To run:
    PYTHONPATH=src pytest tests/test_memory_simulation.py -q -s
"""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path
from textwrap import dedent

import pytest

# The data_sources/, schemas/, and scripts/ directories live at the project
# root, not under src/. Ensure project root is importable regardless of how
# pytest is invoked.
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from data_sources.memory_scenario_loader import load_memory_scenarios  # noqa: E402
from schemas.simulation_case import SimulationCase  # noqa: E402


def _load_simulation_module():
    """
    Load scripts/run_agent_simulation.py as an importable module so we can
    test its functions directly (the file is a script, not a package).
    """
    script_path = PROJECT_ROOT / "scripts" / "run_agent_simulation.py"
    spec = importlib.util.spec_from_file_location(
        "run_agent_simulation_under_test", script_path
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


SIM = _load_simulation_module()


# ---------------------------------------------------------------------------
# Test #1: YAML scenarios load correctly
# ---------------------------------------------------------------------------

def test_default_memory_scenarios_load() -> None:
    """The shipped memory_scenarios.yaml loads into SimulationCase objects
    with all memory fields populated."""
    scenarios = load_memory_scenarios()

    assert len(scenarios) >= 6, "expected at least 6 default memory scenarios"
    for scn in scenarios:
        assert isinstance(scn, SimulationCase)
        assert scn.memory_seed, f"{scn.persona_id} missing memory_seed"
        assert scn.memory_recall_turn is not None, (
            f"{scn.persona_id} missing memory_recall_turn"
        )
        assert scn.expected_recall, f"{scn.persona_id} missing expected_recall"
        assert scn.recall_question, f"{scn.persona_id} missing recall_question"
        assert scn.distractor_turns, f"{scn.persona_id} missing distractor_turns"


def test_multi_detail_scenario_uses_list_expected_recall() -> None:
    """Scenario 6 was upgraded to a true multi-detail recall test, so its
    expected_recall must be a list with both seeded facts."""
    scenarios = load_memory_scenarios()
    multi = next(s for s in scenarios if s.persona_id == "MEM-P06")

    assert isinstance(multi.expected_recall, list)
    expected_lower = [t.lower() for t in multi.expected_recall]
    assert "seoul" in expected_lower
    assert "three months" in expected_lower


def test_load_custom_memory_yaml_roundtrip(tmp_path) -> None:
    """A user-supplied YAML path is honored and parsed into SimulationCases."""
    yaml_text = dedent(
        """
        scenarios:
          - persona_id: "T-P01"
            module_id: "T-M01"
            prompt: "test prompt"
            emotion_before: "calm"
            ai_behavior: "be supportive"
            tone: "warm"
            memory_seed: "I have a cat named Mochi."
            memory_recall_turn: 2
            expected_recall: "Mochi"
            recall_question: "What was my cat's name?"
            qa_check_type: "contains_keyword"
            distractor_turns:
              - "filler 1"
        """
    ).strip()
    custom = tmp_path / "custom_mem.yaml"
    custom.write_text(yaml_text, encoding="utf-8")

    scenarios = load_memory_scenarios(custom)

    assert len(scenarios) == 1
    s = scenarios[0]
    assert s.persona_id == "T-P01"
    assert s.expected_recall == "Mochi"
    assert s.memory_recall_turn == 2


# ---------------------------------------------------------------------------
# Test #2: recall question is injected on the expected turn
# ---------------------------------------------------------------------------

def _make_scenario(**overrides) -> SimulationCase:
    base = dict(
        persona_id="MEM-TEST",
        module_id="MEM-MOD",
        prompt="Walking through the park.",
        emotion_before="calm",
        ai_behavior="listen",
        tone="warm",
        memory_seed="My dog's name is Benny.",
        memory_recall_turn=3,
        expected_recall="Benny",
        recall_question="What was my dog's name?",
        qa_check_type="contains_keyword",
        distractor_turns=[
            "Distractor one.",
            "Distractor two.",
        ],
    )
    base.update(overrides)
    return SimulationCase(**base)


def test_scripted_user_input_seeds_memory_on_turn_zero() -> None:
    scn = _make_scenario()
    text = SIM.scripted_user_input(0, scn)
    assert scn.memory_seed in text
    assert scn.prompt.split(".")[0] in text  # prompt body is included


def test_scripted_user_input_returns_distractors_between_seed_and_recall() -> None:
    scn = _make_scenario()
    assert SIM.scripted_user_input(1, scn) == scn.distractor_turns[0]
    assert SIM.scripted_user_input(2, scn) == scn.distractor_turns[1]


def test_scripted_user_input_injects_recall_question_on_recall_turn() -> None:
    """The recall_question must surface on the configured recall turn."""
    scn = _make_scenario(memory_recall_turn=3)
    assert SIM.scripted_user_input(3, scn) == scn.recall_question


def test_scripted_user_input_injects_recall_question_when_recall_turn_changes() -> None:
    """Sanity check: changing memory_recall_turn moves where the question
    fires."""
    scn = _make_scenario(memory_recall_turn=4, distractor_turns=["a", "b", "c"])
    assert SIM.scripted_user_input(3, scn) == "c"  # last distractor
    assert SIM.scripted_user_input(4, scn) == scn.recall_question


# ---------------------------------------------------------------------------
# Test #3: check_memory_recall hit / miss behavior
# ---------------------------------------------------------------------------

def test_check_memory_recall_hit_single_keyword() -> None:
    scn = _make_scenario(expected_recall="Benny", memory_recall_turn=3)
    result = SIM.check_memory_recall(
        "Yes — you mentioned Benny earlier, your dog.", scn, turn=3
    )

    assert result["is_recall_turn"] is True
    assert result["recall_hit"] == 1
    assert result["matched_terms"] == ["benny"]
    assert result["missing_terms"] == []


def test_check_memory_recall_miss_single_keyword() -> None:
    scn = _make_scenario(expected_recall="Benny", memory_recall_turn=3)
    result = SIM.check_memory_recall(
        "I don't think we've discussed any pets yet.", scn, turn=3
    )

    assert result["is_recall_turn"] is True
    assert result["recall_hit"] == 0
    assert "benny" in result["missing_terms"]


def test_check_memory_recall_returns_none_off_recall_turn() -> None:
    scn = _make_scenario(memory_recall_turn=3)
    result = SIM.check_memory_recall("any response", scn, turn=2)

    assert result["is_recall_turn"] is False
    assert result["recall_hit"] is None


def test_check_memory_recall_skips_when_not_a_memory_scenario() -> None:
    scn = SimulationCase(
        persona_id="NOMEM",
        module_id="X",
        prompt="p",
        emotion_before="calm",
    )
    result = SIM.check_memory_recall("anything", scn, turn=3)
    assert result["is_recall_turn"] is False
    assert result["recall_hit"] is None


def test_check_memory_recall_multi_detail_hit_requires_all_terms() -> None:
    """Multi-detail recall must require *every* listed keyword."""
    scn = _make_scenario(
        expected_recall=["Seoul", "three months"],
        memory_recall_turn=3,
    )
    response = "You mentioned moving from Seoul about three months ago."
    result = SIM.check_memory_recall(response, scn, turn=3)

    assert result["recall_hit"] == 1
    assert set(result["matched_terms"]) == {"seoul", "three months"}
    assert result["missing_terms"] == []


def test_check_memory_recall_multi_detail_miss_when_any_term_absent() -> None:
    """Reproducing only one of two seeded facts must be a MISS."""
    scn = _make_scenario(
        expected_recall=["Seoul", "three months"],
        memory_recall_turn=3,
    )
    response = "You mentioned moving from Seoul, right?"
    result = SIM.check_memory_recall(response, scn, turn=3)

    assert result["recall_hit"] == 0
    assert "three months" in result["missing_terms"]
    assert "seoul" in result["matched_terms"]


# ---------------------------------------------------------------------------
# Test #4: check_continuity catches obvious breaks
# ---------------------------------------------------------------------------

@pytest.mark.parametrize(
    "response",
    [
        "I don't know what you're referring to.",
        "We haven't discussed that before.",
        "I'm not sure what you mean by earlier.",
        "This is our first conversation, so I have no context.",
        "I don't have any previous context for this chat.",
    ],
)
def test_check_continuity_flags_obvious_breaks(response: str) -> None:
    scn = _make_scenario()
    result = SIM.check_continuity(response, turn=2, scenario=scn)
    assert result["continuity_break"] == 1, (
        f"expected continuity break for response: {response!r}"
    )


def test_check_continuity_passes_normal_response() -> None:
    scn = _make_scenario()
    result = SIM.check_continuity(
        "Of course — you mentioned earlier that Benny is your dog.",
        turn=2,
        scenario=scn,
    )
    assert result["continuity_break"] == 0


def test_check_continuity_skips_turn_zero_and_non_memory_scenarios() -> None:
    scn = _make_scenario()
    # turn 0 is the opening, before memory is seeded — not evaluated
    assert SIM.check_continuity("any text", turn=0, scenario=scn)["continuity_break"] is None

    # non-memory scenarios skip the continuity check entirely
    plain = SimulationCase(
        persona_id="NOMEM",
        module_id="X",
        prompt="p",
        emotion_before="calm",
    )
    assert SIM.check_continuity(
        "I don't know what you're referring to.", turn=2, scenario=plain
    )["continuity_break"] is None