"""
Tests for MessageRouter.
"""

import pytest
from dcos.communication.router import MessageRouter
from dcos.communication.protocol import Message


class TestMessageRouter:
    def test_register_and_route(self):
        router = MessageRouter()
        router.register("weather", "agent_weather")
        msg = Message(sender="user", recipient="", payload="weather forecast")
        target = router.route(msg)
        assert target == "agent_weather"

    def test_fallback(self):
        router = MessageRouter()
        router.set_fallback("default_agent")
        msg = Message(sender="user", recipient="", payload="unknown topic")
        target = router.route(msg)
        assert target == "default_agent"

    def test_unregister(self):
        router = MessageRouter()
        router.register("key", "agent")
        removed = router.unregister("key")
        assert removed == "agent"
        msg = Message(sender="user", recipient="", payload="key data")
        target = router.route(msg)
        assert target is None