tests/runtime/test_memory_frames_multiagent.py

# tests/runtime/test_memory_frames_multiagent.py
"""Multi-agent memory frames paint shapes + colours; single-agent path unchanged."""
from proteus.game.runtime.memory import (
    AgentFrame, MemoryCheckpoint, MemoryTurn, memory_frames,
)

_LEGEND = {5: ".", 1: "A", 2: "B", 3: "#", 14: "F"}
_GRID = (12, 12)
_CHOSEN, _DISTRACT, _PRED, _FOOD, _BG = 1, 9, 2, 14, 5


def _ck(turn):
    return MemoryCheckpoint(
        model="m", scenario="s", difficulty="easy", created_at="x",
        outcome="survived", transparent_prompt="p", memory_turns=[turn],
    )


def test_multiagent_paints_agents_predator_resources():
    turn = MemoryTurn(
        turn_idx=1, frame_ascii="", action="up", focal_pos=(0, 0), predator_pos=(0, 0),
        agents=[
            AgentFrame(id="a0", kind="agent", pos=(2, 2), size=2, is_chosen=True),
            AgentFrame(id="a1", kind="agent", pos=(6, 6), size=2),
            AgentFrame(id="a2", kind="agent", pos=(1, 8), size=2, alive=False),
            AgentFrame(id="predator", kind="predator", pos=(8, 1), size=3, facing="right"),
        ],
        resources=[(5, 5)], events=["a2 eaten"],
    )
    grid = memory_frames(_ck(turn), legend=_LEGEND, grid_size=_GRID)[0]["grid"]
    # chosen 2x2 at (2,2)
    assert grid[2][2] == _CHOSEN and grid[3][3] == _CHOSEN
    # distractor 2x2 at (6,6) in blue
    assert grid[6][6] == _DISTRACT
    # eaten agent NOT painted
    assert grid[8][1] == _BG
    # resource
    assert grid[5][5] == _FOOD
    # predator ㄷ at (8,1), mouth EAST => the two right-middle cells are background
    assert grid[1][8] == _PRED and grid[3][10] == _PRED   # corners solid
    assert grid[2][9] == _BG and grid[2][10] == _BG        # mouth cells transparent


def test_multiagent_frame_carries_events():
    turn = MemoryTurn(turn_idx=1, frame_ascii="", action="up", focal_pos=(0, 0),
                      predator_pos=(0, 0),
                      agents=[AgentFrame(id="a0", kind="agent", pos=(0, 0), size=2)],
                      events=["hello"])
    out = memory_frames(_ck(turn), legend=_LEGEND, grid_size=_GRID)[0]
    assert out["events"] == ["hello"]


def test_single_agent_path_unchanged():
    # No agents -> legacy reconstruction (prose frame), paints predator 5-block etc.
    turn = MemoryTurn(turn_idx=1, frame_ascii="prose", action="up",
                      focal_pos=(2, 2), predator_pos=(7, 7))
    out = memory_frames(_ck(turn), legend=_LEGEND, grid_size=_GRID)[0]
    assert "grid" in out and out["events"] == []