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."""