"""
Tests for AddressResolver.
"""

import pytest
from dcos.communication.resolver import AddressResolver


class TestAddressResolver:
    def test_resolve_by_name(self):
        resolver = AddressResolver()
        resolver.register("planner", "agent_001", role="planner")
        assert resolver.resolve_by_name("planner") == "agent_001"

    def test_resolve_by_role(self):
        resolver = AddressResolver()
        resolver.register("planner1", "agent_001", role="planner")
        resolver.register("planner2", "agent_002", role="planner")
        planners = resolver.resolve_by_role("planner")
        assert len(planners) == 2

    def test_unregister(self):
        resolver = AddressResolver()
        resolver.register("temp", "agent_temp")
        removed = resolver.unregister("temp")
        assert removed == "agent_temp"
        assert resolver.resolve_by_name("temp") is None