# Path: azmq/multiplexer.py
# class Multiplexer(ClosableAsyncObject):
#     def on_open(self):
#         super().on_open()
#         self._sockets = set()
# 
#     def add_socket(self, socket):
#         """
#         Add a socket to the multiplexer.
# 
#         :param socket: The socket. If it was added already, it won't be added a
#             second time.
#         """
#         if socket not in self._sockets:
#             self._sockets.add(socket)
#             socket.on_closed.connect(self.remove_socket)
# 
#     def remove_socket(self, socket):
#         """
#         Remove a socket from the multiplexer.
# 
#         :param socket: The socket. If it was removed already or if it wasn't
#             added, the call does nothing.
#         """
#         if socket in self._sockets:
#             socket.on_closed.disconnect(self.remove_socket)
#             self._sockets.remove(socket)
# 
#     @cancel_on_closing
#     async def recv_multipart(self):
#         """
#         Read from all the associated sockets.
# 
#         :returns: A list of tuples (socket, frames) for each socket that
#             returned a result.
#         """
#         if not self._sockets:
#             return []
# 
#         results = []
# 
#         async def recv_and_store(socket):
#             frames = await socket.recv_multipart()
#             results.append((socket, frames))
# 
#         tasks = [
#             asyncio.ensure_future(recv_and_store(socket), loop=self.loop)
#             for socket in self._sockets
#         ]
# 
#         try:
#             await asyncio.wait(
#                 tasks,
#                 return_when=asyncio.FIRST_COMPLETED,
#                 loop=self.loop,
#             )
#         finally:
#             for task in tasks:
#                 task.cancel()
# 
#         return results


import pytest
from mock import MagicMock
from azmq.multiplexer import Multiplexer
"""
Unit tests for the multiplexer.
"""





@pytest.mark.asyncio
async def test_add_socket_twice(event_loop):
    socket = MagicMock()
    multiplexer = Multiplexer(loop=event_loop)
    multiplexer.add_socket(socket)
    multiplexer.add_socket(socket)

    assert multiplexer._sockets == {socket}


@pytest.mark.asyncio
async def test_remove_unadded_socket(event_loop):
    socket = MagicMock()
    multiplexer = Multiplexer(loop=event_loop)
    multiplexer.remove_socket(socket)

    assert multiplexer._sockets == set()


@pytest.mark.asyncio
async def test_recv_multipart_no_sockets(event_loop):

# Path: azmq/multiplexer.py
# class Multiplexer(ClosableAsyncObject):
#     def on_open(self):
#         super().on_open()
#         self._sockets = set()
# 
#     def add_socket(self, socket):
#         """
#         Add a socket to the multiplexer.
# 
#         :param socket: The socket. If it was added already, it won't be added a
#             second time.
#         """
#         if socket not in self._sockets:
#             self._sockets.add(socket)
#             socket.on_closed.connect(self.remove_socket)
# 
#     def remove_socket(self, socket):
#         """
#         Remove a socket from the multiplexer.
# 
#         :param socket: The socket. If it was removed already or if it wasn't
#             added, the call does nothing.
#         """
#         if socket in self._sockets:
#             socket.on_closed.disconnect(self.remove_socket)
#             self._sockets.remove(socket)
# 
#     @cancel_on_closing
#     async def recv_multipart(self):
#         """
#         Read from all the associated sockets.
# 
#         :returns: A list of tuples (socket, frames) for each socket that
#             returned a result.
#         """
#         if not self._sockets:
#             return []
# 
#         results = []
# 
#         async def recv_and_store(socket):
#             frames = await socket.recv_multipart()
#             results.append((socket, frames))
# 
#         tasks = [
#             asyncio.ensure_future(recv_and_store(socket), loop=self.loop)
#             for socket in self._sockets
#         ]
# 
#         try:
#             await asyncio.wait(
#                 tasks,
#                 return_when=asyncio.FIRST_COMPLETED,
#                 loop=self.loop,
#             )
#         finally:
#             for task in tasks:
#                 task.cancel()
# 
#         return results


# Path: tests/unit/test_multiplexer.py
import pytest
from mock import MagicMock
from azmq.multiplexer import Multiplexer

"""
Unit tests for the multiplexer.
"""





@pytest.mark.asyncio
async def test_add_socket_twice(event_loop):
    socket = MagicMock()
    multiplexer = Multiplexer(loop=event_loop)
    multiplexer.add_socket(socket)
    multiplexer.add_socket(socket)

    assert multiplexer._sockets == {socket}


@pytest.mark.asyncio
async def test_remove_unadded_socket(event_loop):
    socket = MagicMock()
    multiplexer = Multiplexer(loop=event_loop)
    multiplexer.remove_socket(socket)

    assert multiplexer._sockets == set()


@pytest.mark.asyncio
async def test_recv_multipart_no_sockets(event_loop):

    multiplexer = Multiplexer(loop=event_loop)