AgentEvalTool/tests/unit/mock_channel.py

72 lines
2.3 KiB
Python

"""Mock EvalChannel implementation for unit tests.
Simulates send/poll_reply with configurable latency, failure modes, and
cancellation hooks. Does not touch the network.
"""
import asyncio
from typing import Any, Optional
from agenteval.channels.base import ChannelHealth, EvalChannel, Reply, SendResult
class MockChannel(EvalChannel):
"""Configurable in-memory channel for engine tests."""
def __init__(
self,
reply_delay: float = 0.0,
send_ok: bool = True,
missing_reply: bool = False,
reply_text: Optional[str] = None,
raise_on_send: Optional[Exception] = None,
raise_on_poll: Optional[Exception] = None,
):
self.reply_delay = reply_delay
self.send_ok = send_ok
self.missing_reply = missing_reply
self.reply_text = reply_text
self.raise_on_send = raise_on_send
self.raise_on_poll = raise_on_poll
self.sent: list[str] = []
self.send_calls = 0
self.poll_calls = 0
self._msg_counter = 0
async def health_check(self) -> ChannelHealth:
return ChannelHealth(ok=True, message="mock")
async def _send(self, content: str, **kwargs: Any) -> SendResult:
self.send_calls += 1
self.sent.append(content)
if self.raise_on_send:
raise self.raise_on_send
if not self.send_ok:
return SendResult(ok=False, error="mock send failure")
self._msg_counter += 1
return SendResult(ok=True, question_msg_id=f"q-{self._msg_counter}")
async def _poll_reply(
self,
question_msg_id: str,
timeout: float = 30.0,
poll_interval: float = 1.0,
) -> Optional[Reply]:
self.poll_calls += 1
if self.raise_on_poll:
raise self.raise_on_poll
if self.missing_reply:
# Simulate a slow target: sleep past the timeout so the engine
# observes a poll timeout.
await asyncio.sleep(timeout + 0.05)
return None
if self.reply_delay:
await asyncio.sleep(self.reply_delay)
text = self.reply_text if self.reply_text is not None else f"echo: {question_msg_id}"
return Reply(
question_msg_id=question_msg_id,
content=text,
raw_message={"msgBody": {"content": text}},
)