GH-98388: Add tests for happy eyeballs and its internal workings

Lib/test/test_asyncio/test_base_events.py

@@ -145,6 +145,27 @@ def test_ipaddr_info_no_inet_pton(self, m_socket):
  socket.SOCK_STREAM,
  socket.IPPROTO_TCP))
 
+ def test_interleave_addrinfos(self):
+ SIX_A = (socket.AF_INET6, 0, 0, '', ('2001:db8::1', 1))
+ SIX_B = (socket.AF_INET6, 0, 0, '', ('2001:db8::2', 2))
+ SIX_C = (socket.AF_INET6, 0, 0, '', ('2001:db8::3', 3))
+ SIX_D = (socket.AF_INET6, 0, 0, '', ('2001:db8::4', 4))
+ FOUR_A = (socket.AF_INET, 0, 0, '', ('192.0.2.1', 5))
+ FOUR_B = (socket.AF_INET, 0, 0, '', ('192.0.2.2', 6))
+ FOUR_C = (socket.AF_INET, 0, 0, '', ('192.0.2.3', 7))
+ FOUR_D = (socket.AF_INET, 0, 0, '', ('192.0.2.4', 8))
+
+ addrinfos = [SIX_A, SIX_B, SIX_C, SIX_D, FOUR_A, FOUR_B, FOUR_C, FOUR_D]
+ expected = [SIX_A, FOUR_A, SIX_B, FOUR_B, SIX_C, FOUR_C, SIX_D, FOUR_D]
+
+ self.assertEqual(expected, base_events._interleave_addrinfos(addrinfos))
+
+ expected_fafc_2 = [SIX_A, SIX_B, FOUR_A, SIX_C, FOUR_B, SIX_D, FOUR_C, FOUR_D]
+ self.assertEqual(
+ expected_fafc_2,
+ base_events._interleave_addrinfos(addrinfos, first_address_family_count=2),
+ )
+
 
 class BaseEventLoopTests(test_utils.TestCase):
 
@@ -1426,6 +1447,65 @@ def getaddrinfo_task(*args, **kwds):
  self.assertRaises(
  OSError, self.loop.run_until_complete, coro)
 
+ @unittest.skipUnless(socket_helper.IPV6_ENABLED, 'no IPv6 support')
+ @patch_socket
+ def test_create_connection_happy_eyeballs(self, m_socket):
+
+ class MyProto(asyncio.Protocol):
+ pass
+
+ async def getaddrinfo(*args, **kw):
+ return [(socket.AF_INET6, 0, 0, '', ('2001:db8::1', 1)),
+ (socket.AF_INET, 0, 0, '', ('192.0.2.1', 5))]
+
+ async def sock_connect(sock, address):
+ if address[0] == '2001:db8::1':
+ await asyncio.sleep(1)
+ sock.connect(address)
+
+ self.loop._add_reader = mock.Mock()
+ self.loop._add_writer = mock.Mock()
+ self.loop.getaddrinfo = getaddrinfo
+ self.loop.sock_connect = sock_connect
+
+ coro = self.loop.create_connection(MyProto, 'example.com', 80, happy_eyeballs_delay=0.3)
+ transport, protocol = self.loop.run_until_complete(coro)
+ try:
+ sock = transport._sock
+ sock.connect.assert_called_with(('192.0.2.1', 5))
+ finally:
+ transport.close()
+ test_utils.run_briefly(self.loop) # allow transport to close
+
+ @patch_socket
+ def test_create_connection_happy_eyeballs_ipv4_only(self, m_socket):
+
+ class MyProto(asyncio.Protocol):
+ pass
+
+ async def getaddrinfo(*args, **kw):
+ return [(socket.AF_INET, 0, 0, '', ('192.0.2.1', 5)),
+ (socket.AF_INET, 0, 0, '', ('192.0.2.2', 6))]
+
+ async def sock_connect(sock, address):
+ if address[0] == '192.0.2.1':
+ await asyncio.sleep(1)
+ sock.connect(address)
+
+ self.loop._add_reader = mock.Mock()
+ self.loop._add_writer = mock.Mock()
+ self.loop.getaddrinfo = getaddrinfo
+ self.loop.sock_connect = sock_connect
+
+ coro = self.loop.create_connection(MyProto, 'example.com', 80, happy_eyeballs_delay=0.3)
+ transport, protocol = self.loop.run_until_complete(coro)
+ try:
+ sock = transport._sock
+ sock.connect.assert_called_with(('192.0.2.2', 6))
+ finally:
+ transport.close()
+ test_utils.run_briefly(self.loop) # allow transport to close
+
  @patch_socket
  def test_create_connection_bluetooth(self, m_socket):
  # See http://bugs.python.org/issue27136, fallback to getaddrinfo when

Lib/test/test_asyncio/test_staggered.py

@@ -0,0 +1,115 @@
+import asyncio
+import functools
+import unittest
+from asyncio.staggered import staggered_race
+
+
+# To prevent a warning "test altered the execution environment"
+def tearDownModule():
+ asyncio.set_event_loop_policy(None)
+
+
+class TestStaggered(unittest.IsolatedAsyncioTestCase):
+ @staticmethod
+ async def waiting_coroutine(return_value, wait_seconds, success):
+ await asyncio.sleep(wait_seconds)
+ if success:
+ return return_value
+ raise RuntimeError(str(return_value))
+
+ def get_waiting_coroutine_factory(self, return_value, wait_seconds, success):
+ return functools.partial(self.waiting_coroutine, return_value, wait_seconds, success)
+
+ async def test_single_success(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (self.get_waiting_coroutine_factory(0, 0.1, True),),
+ 0.1,
+ )
+ self.assertEqual(winner_result, 0)
+ self.assertEqual(winner_idx, 0)
+ self.assertEqual(len(exceptions), 1)
+ self.assertIsNone(exceptions[0])
+
+ async def test_single_fail(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (self.get_waiting_coroutine_factory(0, 0.1, False),),
+ 0.1,
+ )
+ self.assertEqual(winner_result, None)
+ self.assertEqual(winner_idx, None)
+ self.assertEqual(len(exceptions), 1)
+ self.assertIsInstance(exceptions[0], RuntimeError)
+
+ async def test_first_win(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (
+ self.get_waiting_coroutine_factory(0, 0.2, True),
+ self.get_waiting_coroutine_factory(1, 0.2, True),
+ ),
+ 0.1,
+ )
+ self.assertEqual(winner_result, 0)
+ self.assertEqual(winner_idx, 0)
+ self.assertEqual(len(exceptions), 2)
+ self.assertIsNone(exceptions[0])
+ self.assertIsInstance(exceptions[1], asyncio.CancelledError)
+
+ async def test_second_win(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (
+ self.get_waiting_coroutine_factory(0, 0.3, True),
+ self.get_waiting_coroutine_factory(1, 0.1, True),
+ ),
+ 0.1,
+ )
+ self.assertEqual(winner_result, 1)
+ self.assertEqual(winner_idx, 1)
+ self.assertEqual(len(exceptions), 2)
+ self.assertIsInstance(exceptions[0], asyncio.CancelledError)
+ self.assertIsNone(exceptions[1])
+
+ async def test_first_fail(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (
+ self.get_waiting_coroutine_factory(0, 0.2, False),
+ self.get_waiting_coroutine_factory(1, 0.2, True),
+ ),
+ 0.1,
+ )
+ self.assertEqual(winner_result, 1)
+ self.assertEqual(winner_idx, 1)
+ self.assertEqual(len(exceptions), 2)
+ self.assertIsInstance(exceptions[0], RuntimeError)
+ self.assertIsNone(exceptions[1])
+
+ async def test_second_fail(self):
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (
+ self.get_waiting_coroutine_factory(0, 0.2, True),
+ self.get_waiting_coroutine_factory(1, 0, False),
+ ),
+ 0.1,
+ )
+ self.assertEqual(winner_result, 0)
+ self.assertEqual(winner_idx, 0)
+ self.assertEqual(len(exceptions), 2)
+ self.assertIsNone(exceptions[0])
+ self.assertIsInstance(exceptions[1], RuntimeError)
+
+ async def test_simultaneous_success_fail(self):
+ # There's a potential race condition here:
+ # https://github.com/python/cpython/issues/86296
+ # As with any race condition, it can be difficult to reproduce.
+ # This test may not fail every time.
+ for i in range(201):
+ time_unit = 0.0001 * i
+ winner_result, winner_idx, exceptions = await staggered_race(
+ (
+ self.get_waiting_coroutine_factory(0, time_unit*2, True),
+ self.get_waiting_coroutine_factory(1, time_unit, False),
+ self.get_waiting_coroutine_factory(2, 0.05, True)
+ ),
+ time_unit,
+ )
+ self.assertEqual(winner_result, 0)
+ self.assertEqual(winner_idx, 0)