Add unused port helpers for UDP (#99)

* Add unused port helpers for UDP

Extends the unused_tcp_port and unused_tcp_port_factory mechanisms for
UDP ports.

* Update pytest_asyncio/plugin.py

* Add changenote

Co-authored-by: Andrew Svetlov <[email protected]>

Author: dbuse

README.rst

@@ -109,6 +109,11 @@ when several unused TCP ports are required in a test.
         port1, port2 = unused_tcp_port_factory(), unused_tcp_port_factory()
         ...
 
+``unused_udp_port`` and ``unused_udp_port_factory``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Work just like their TCP counterparts but return unused UDP ports.
+
+
 Async fixtures
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Asynchronous fixtures are defined just like ordinary pytest fixtures, except they should be coroutines or asynchronous generators.
@@ -166,7 +171,7 @@ Note about unittest
 -------------------
 
 Test classes subclassing the standard `unittest <https://docs.python.org/3/library/unittest.html>`__ library are not supported, users
-are recommended to use `unitest.IsolatedAsyncioTestCase <https://docs.python.org/3/library/unittest.html#unittest.IsolatedAsyncioTestCase>`__ 
+are recommended to use `unitest.IsolatedAsyncioTestCase <https://docs.python.org/3/library/unittest.html#unittest.IsolatedAsyncioTestCase>`__
 or an async framework such as `asynctest <https://asynctest.readthedocs.io/en/latest>`__.
 
 Changelog
@@ -177,6 +182,7 @@ Changelog
 - Drop support for Python 3.6
 - Fixed an issue when pytest-asyncio was used in combination with `flaky` or inherited asynchronous Hypothesis tests. `#178 <https://github.com/pytest-dev/pytest-asyncio/issues/178>`_ `#231 <https://github.com/pytest-dev/pytest-asyncio/issues/231>`_
 - Added `flaky <https://pypi.org/project/flaky/>`_ to test dependencies
+- Added ``unused_udp_port`` and ``unused_udp_port_factory`` fixtures (similar to ``unused_tcp_port`` and ``unused_tcp_port_factory`` counterparts. `#99 <https://github.com/pytest-dev/pytest-asyncio/issues/99>`_
 
 0.16.0 (2021-10-16)
 ~~~~~~~~~~~~~~~~~~~

pytest_asyncio/plugin.py

@@ -224,16 +224,21 @@ def event_loop(request):
     loop.close()
 
 
-def _unused_tcp_port():
-    """Find an unused localhost TCP port from 1024-65535 and return it."""
-    with contextlib.closing(socket.socket()) as sock:
+def _unused_port(socket_type):
+    """Find an unused localhost port from 1024-65535 and return it."""
+    with contextlib.closing(socket.socket(type=socket_type)) as sock:
         sock.bind(("127.0.0.1", 0))
         return sock.getsockname()[1]
 
 
 @pytest.fixture
 def unused_tcp_port():
-    return _unused_tcp_port()
+    return _unused_port(socket.SOCK_STREAM)
+
+
+@pytest.fixture
+def unused_udp_port():
+    return _unused_port(socket.SOCK_DGRAM)
 
 
 @pytest.fixture(scope="session")
@@ -243,10 +248,29 @@ def unused_tcp_port_factory():
 
     def factory():
         """Return an unused port."""
-        port = _unused_tcp_port()
+        port = _unused_port(socket.SOCK_STREAM)
+
+        while port in produced:
+            port = _unused_port(socket.SOCK_STREAM)
+
+        produced.add(port)
+
+        return port
+
+    return factory
+
+
+@pytest.fixture(scope="session")
+def unused_udp_port_factory():
+    """A factory function, producing different unused UDP ports."""
+    produced = set()
+
+    def factory():
+        """Return an unused port."""
+        port = _unused_port(socket.SOCK_DGRAM)
 
         while port in produced:
-            port = _unused_tcp_port()
+            port = _unused_port(socket.SOCK_DGRAM)
 
         produced.add(port)
 

tests/test_simple.py

@@ -51,6 +51,33 @@ async def closer(_, writer):
     await server1.wait_closed()
 
 
+@pytest.mark.asyncio
+async def test_unused_udp_port_fixture(unused_udp_port, event_loop):
+    """Test the unused TCP port fixture."""
+
+    class Closer:
+        def connection_made(self, transport):
+            pass
+
+        def connection_lost(self, *arg, **kwd):
+            pass
+
+    transport1, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=("127.0.0.1", unused_udp_port),
+        reuse_port=False,
+    )
+
+    with pytest.raises(IOError):
+        await event_loop.create_datagram_endpoint(
+            Closer,
+            local_addr=("127.0.0.1", unused_udp_port),
+            reuse_port=False,
+        )
+
+    transport1.abort()
+
+
 @pytest.mark.asyncio
 async def test_unused_port_factory_fixture(unused_tcp_port_factory, event_loop):
     """Test the unused TCP port factory fixture."""
@@ -80,11 +107,57 @@ async def closer(_, writer):
     await server3.wait_closed()
 
 
+@pytest.mark.asyncio
+async def test_unused_udp_port_factory_fixture(unused_udp_port_factory, event_loop):
+    """Test the unused UDP port factory fixture."""
+
+    class Closer:
+        def connection_made(self, transport):
+            pass
+
+        def connection_lost(self, *arg, **kwd):
+            pass
+
+    port1, port2, port3 = (
+        unused_udp_port_factory(),
+        unused_udp_port_factory(),
+        unused_udp_port_factory(),
+    )
+
+    transport1, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=("127.0.0.1", port1),
+        reuse_port=False,
+    )
+    transport2, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=("127.0.0.1", port2),
+        reuse_port=False,
+    )
+    transport3, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=("127.0.0.1", port3),
+        reuse_port=False,
+    )
+
+    for port in port1, port2, port3:
+        with pytest.raises(IOError):
+            await event_loop.create_datagram_endpoint(
+                Closer,
+                local_addr=("127.0.0.1", port),
+                reuse_port=False,
+            )
+
+    transport1.abort()
+    transport2.abort()
+    transport3.abort()
+
+
 def test_unused_port_factory_duplicate(unused_tcp_port_factory, monkeypatch):
     """Test correct avoidance of duplicate ports."""
     counter = 0
 
-    def mock_unused_tcp_port():
+    def mock_unused_tcp_port(_ignored):
         """Force some duplicate ports."""
         nonlocal counter
         counter += 1
@@ -93,12 +166,31 @@ def mock_unused_tcp_port():
         else:
             return 10000 + counter
 
-    monkeypatch.setattr(pytest_asyncio.plugin, "_unused_tcp_port", mock_unused_tcp_port)
+    monkeypatch.setattr(pytest_asyncio.plugin, "_unused_port", mock_unused_tcp_port)
 
     assert unused_tcp_port_factory() == 10000
     assert unused_tcp_port_factory() > 10000
 
 
+def test_unused_udp_port_factory_duplicate(unused_udp_port_factory, monkeypatch):
+    """Test correct avoidance of duplicate UDP ports."""
+    counter = 0
+
+    def mock_unused_udp_port(_ignored):
+        """Force some duplicate ports."""
+        nonlocal counter
+        counter += 1
+        if counter < 5:
+            return 10000
+        else:
+            return 10000 + counter
+
+    monkeypatch.setattr(pytest_asyncio.plugin, "_unused_port", mock_unused_udp_port)
+
+    assert unused_udp_port_factory() == 10000
+    assert unused_udp_port_factory() > 10000
+
+
 class TestMarkerInClassBasedTests:
     """Test that asyncio marked functions work for methods of test classes."""