import asyncio
import pytest
from protocols.mcp_protocol import MCPProtocol
from protocols.base_protocol import Message


@pytest.mark.asyncio
async def test_mcp_send_receive():
    mcp = MCPProtocol()
    msg = Message(sender="agent_a", recipient="agent_b", msg_type="task_complete", job_id="j1")
    await mcp.send(msg)
    received = await mcp.receive("agent_b", timeout=1.0)
    assert received is not None
    assert received.sender == "agent_a"
    assert received.msg_type == "task_complete"


@pytest.mark.asyncio
async def test_mcp_receive_timeout():
    mcp = MCPProtocol()
    result = await mcp.receive("nobody", timeout=0.1)
    assert result is None


@pytest.mark.asyncio
async def test_mcp_broadcast():
    mcp = MCPProtocol()
    msg = Message(sender="orchestrator", recipient="", msg_type="start", job_id="j2")
    await mcp.broadcast(msg, recipients=["agent_a", "agent_b"])
    a = await mcp.receive("agent_a", timeout=1.0)
    b = await mcp.receive("agent_b", timeout=1.0)
    assert a is not None and b is not None