Add unused port helpers for UDP

README.rst

@@ -107,6 +107,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.

pytest_asyncio/plugin.py

@@ -174,16 +174,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=socket.SOCK_STREAM):
+    """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()
+
+
+@pytest.fixture
+def unused_udp_port():
+    return _unused_port(socket.SOCK_DGRAM)
 
 
 @pytest.fixture
@@ -193,10 +198,28 @@ def unused_tcp_port_factory():
 
     def factory():
         """Return an unused port."""
-        port = _unused_tcp_port()
+        port = _unused_port()
+
+        while port in produced:
+            port = _unused_port()
+
+        produced.add(port)
+
+        return port
+    return factory
+
+
+@pytest.fixture
+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

@@ -57,6 +57,35 @@ 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,
+        reuse_address=False,
+    )
+
+    with pytest.raises(IOError):
+        await event_loop.create_datagram_endpoint(
+            Closer,
+            local_addr=('127.0.0.1', unused_udp_port),
+            reuse_port=False,
+            reuse_address=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."""
@@ -91,6 +120,54 @@ 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,
+        reuse_address=False,
+    )
+    transport2, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=('127.0.0.1', port2),
+        reuse_port=False,
+        reuse_address=False,
+    )
+    transport3, _ = await event_loop.create_datagram_endpoint(
+        Closer,
+        local_addr=('127.0.0.1', port3),
+        reuse_port=False,
+        reuse_address=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,
+                reuse_address=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
@@ -104,13 +181,34 @@ def mock_unused_tcp_port():
         else:
             return 10000 + counter
 
-    monkeypatch.setattr(pytest_asyncio.plugin, '_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 Test:
     """Test that asyncio marked functions work in test methods."""