"""Unit tests for bannerlord agents — King, Vassals, Companions.""" import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest from bannerlord.agents.companions import ( CaravanCompanion, LogisticsCompanion, ScoutCompanion, ) from bannerlord.agents.king import KingAgent from bannerlord.agents.vassals import DiplomacyVassal, EconomyVassal, WarVassal from bannerlord.gabs_client import GABSClient, GABSUnavailable from bannerlord.ledger import Ledger from bannerlord.models import ( KingSubgoal, ResultMessage, SubgoalMessage, TaskMessage, VictoryCondition, ) # ── Helpers ─────────────────────────────────────────────────────────────────── def _mock_gabs(state: dict | None = None) -> GABSClient: """Return a disconnected GABS stub that returns *state* from get_state.""" gabs = MagicMock(spec=GABSClient) gabs.connected = False if state is not None: gabs.get_state = AsyncMock(return_value=state) else: gabs.get_state = AsyncMock(side_effect=GABSUnavailable("no game")) gabs.call = AsyncMock(return_value={}) gabs.recruit_troops = AsyncMock(return_value={"recruited": 10}) gabs.move_party = AsyncMock(return_value={"moving": True}) return gabs def _mock_ledger(tmp_path) -> Ledger: ledger = Ledger(db_path=tmp_path / "ledger.db") ledger.initialize() return ledger # ── King agent ──────────────────────────────────────────────────────────────── class TestKingAgent: async def test_victory_detected(self, tmp_path): """Campaign stops immediately when victory condition is met.""" gabs = _mock_gabs({"player_title": "King", "territory_control_pct": 55.0}) ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0) victory = await king.run_campaign(max_ticks=10) assert victory.achieved async def test_max_ticks_respected(self, tmp_path): """Campaign stops after max_ticks when victory not yet achieved.""" gabs = _mock_gabs({"player_title": "Lord", "territory_control_pct": 10.0}) ledger = _mock_ledger(tmp_path) # Patch LLM to return a valid subgoal without calling Ollama king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0) with patch.object(king, "_decide", AsyncMock(return_value=KingSubgoal(token="RECRUIT"))): victory = await king.run_campaign(max_ticks=3) assert not victory.achieved assert king._tick == 3 async def test_llm_failure_falls_back_to_recruit(self, tmp_path): """If LLM fails, King defaults to RECRUIT subgoal.""" gabs = _mock_gabs({"player_title": "Lord", "territory_control_pct": 5.0}) ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger, tick_interval=0) with patch.object(king, "_llm_decide", side_effect=RuntimeError("Ollama down")): subgoal = await king._decide({}) assert subgoal.token == "RECRUIT" async def test_subgoal_broadcast_to_all_vassals(self, tmp_path): """King broadcasts subgoal to all three vassals.""" gabs = _mock_gabs({}) ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger) subgoal = KingSubgoal(token="EXPAND_TERRITORY", target="Epicrotea") await king._broadcast_subgoal(subgoal) messages = [] while not king.subgoal_queue.empty(): messages.append(king.subgoal_queue.get_nowait()) assert len(messages) == 3 recipients = {m.to_agent for m in messages} assert recipients == {"war_vassal", "economy_vassal", "diplomacy_vassal"} async def test_gabs_unavailable_uses_empty_state(self, tmp_path): """King handles GABS being offline gracefully.""" gabs = _mock_gabs() # raises GABSUnavailable ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger) state = await king._fetch_state() assert state == {} def test_evaluate_victory_king_with_majority(self, tmp_path): gabs = _mock_gabs() ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger) v = king._evaluate_victory({"player_title": "King", "territory_control_pct": 60.0}) assert v.achieved def test_evaluate_victory_not_king(self, tmp_path): gabs = _mock_gabs() ledger = _mock_ledger(tmp_path) king = KingAgent(gabs_client=gabs, ledger=ledger) v = king._evaluate_victory({"player_title": "Lord", "territory_control_pct": 80.0}) assert not v.achieved # ── Vassals ─────────────────────────────────────────────────────────────────── class TestWarVassal: async def test_expand_territory_emits_move_task(self): gabs = _mock_gabs({"territory_delta": 1.0, "army_strength_ratio": 1.5}) queue = asyncio.Queue() vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="EXPAND_TERRITORY", target="Seonon") await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "move_party" assert task.args["destination"] == "Seonon" async def test_recruit_emits_recruit_task(self): gabs = _mock_gabs({}) queue = asyncio.Queue() vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="RECRUIT", quantity=15) await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "recruit_troop" assert task.args["quantity"] == 15 async def test_irrelevant_token_emits_no_task(self): gabs = _mock_gabs({}) queue = asyncio.Queue() vassal = WarVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="ALLY") await vassal._tick(subgoal) assert vassal.task_queue.empty() class TestEconomyVassal: async def test_fortify_emits_build_task(self): gabs = _mock_gabs({"daily_income": 200.0}) queue = asyncio.Queue() vassal = EconomyVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="FORTIFY", target="Epicrotea") await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "build_project" assert task.args["settlement"] == "Epicrotea" async def test_trade_emits_assess_prices(self): gabs = _mock_gabs({}) queue = asyncio.Queue() vassal = EconomyVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="TRADE", target="Pravend") await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "assess_prices" class TestDiplomacyVassal: async def test_ally_emits_track_lord(self): gabs = _mock_gabs({"allies_count": 1}) queue = asyncio.Queue() vassal = DiplomacyVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="ALLY", target="Derthert") await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "track_lord" assert task.args["name"] == "Derthert" async def test_spy_emits_assess_garrison(self): gabs = _mock_gabs({}) queue = asyncio.Queue() vassal = DiplomacyVassal(gabs_client=gabs, subgoal_queue=queue) subgoal = KingSubgoal(token="SPY", target="Marunath") await vassal._tick(subgoal) task: TaskMessage = vassal.task_queue.get_nowait() assert task.primitive == "assess_garrison" assert task.args["settlement"] == "Marunath" # ── Companions ──────────────────────────────────────────────────────────────── class TestLogisticsCompanion: async def test_recruit_troop(self): gabs = _mock_gabs() gabs.recruit_troops = AsyncMock(return_value={"recruited": 10, "type": "infantry"}) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = LogisticsCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="war_vassal", to_agent="logistics_companion", primitive="recruit_troop", args={"troop_type": "infantry", "quantity": 10}, ) result = await comp._execute(task) assert result.success is True assert result.outcome["recruited"] == 10 async def test_unknown_primitive_fails_gracefully(self): gabs = _mock_gabs() q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = LogisticsCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="war_vassal", to_agent="logistics_companion", primitive="launch_nukes", args={}, ) result = await comp._execute(task) assert result.success is False assert "Unknown primitive" in result.outcome["error"] async def test_gabs_unavailable_returns_failure(self): gabs = _mock_gabs() gabs.recruit_troops = AsyncMock(side_effect=GABSUnavailable("offline")) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = LogisticsCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="war_vassal", to_agent="logistics_companion", primitive="recruit_troop", args={"troop_type": "infantry", "quantity": 5}, ) result = await comp._execute(task) assert result.success is False class TestCaravanCompanion: async def test_assess_prices(self): gabs = _mock_gabs() gabs.call = AsyncMock(return_value={"grain": 12, "linen": 45}) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = CaravanCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="economy_vassal", to_agent="caravan_companion", primitive="assess_prices", args={"town": "Pravend"}, ) result = await comp._execute(task) assert result.success is True async def test_abandon_route(self): gabs = _mock_gabs() gabs.call = AsyncMock(return_value={"abandoned": True}) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = CaravanCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="economy_vassal", to_agent="caravan_companion", primitive="abandon_route", args={}, ) result = await comp._execute(task) assert result.success is True assert result.outcome["abandoned"] is True class TestScoutCompanion: async def test_assess_garrison(self): gabs = _mock_gabs() gabs.call = AsyncMock(return_value={"garrison_size": 120, "settlement": "Marunath"}) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = ScoutCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="diplomacy_vassal", to_agent="scout_companion", primitive="assess_garrison", args={"settlement": "Marunath"}, ) result = await comp._execute(task) assert result.success is True assert result.outcome["garrison_size"] == 120 async def test_report_intel(self): gabs = _mock_gabs() gabs.call = AsyncMock(return_value={"intel": ["Derthert at Epicrotea"]}) q: asyncio.Queue[TaskMessage] = asyncio.Queue() comp = ScoutCompanion(gabs_client=gabs, task_queue=q) task = TaskMessage( from_agent="diplomacy_vassal", to_agent="scout_companion", primitive="report_intel", args={}, ) result = await comp._execute(task) assert result.success is True