"""
Tests for CommunicationNetwork.
"""

import pytest
from dcos.communication.network import CommunicationNetwork
from dcos.communication.protocol import Message


class TestCommunicationNetwork:
    def test_connect_and_deliver(self):
        net = CommunicationNetwork()
        net.connect("a", "b")
        msg = Message(sender="a", recipient="b", payload="test")
        assert net.deliver(msg) is True

    def test_disconnect(self):
        net = CommunicationNetwork()
        net.connect("a", "b")
        net.disconnect("a", "b")
        msg = Message(sender="a", recipient="b", payload="test")
        assert net.deliver(msg) is False

    def test_broadcast(self):
        net = CommunicationNetwork()
        net.connect("a", "b")
        net.connect("a", "c")
        count = net.broadcast("a", "broadcast message")
        assert count == 2

    def test_get_peers(self):
        net = CommunicationNetwork()
        net.connect("x", "y")
        net.connect("x", "z")
        peers = net.get_peers("x")
        assert len(peers) == 2