From 6dc8f89ff99a38c8ecaf8045c9afbe683d6f2c6e Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 9 Dec 2023 10:52:16 -1000
Subject: [PATCH] feat: add coverage for multi ipv6 (#10)

---
 tests/test_impl.py | 115 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 115 insertions(+)

diff --git a/tests/test_impl.py b/tests/test_impl.py
index 72d8346..db79d74 100644
--- a/tests/test_impl.py
+++ b/tests/test_impl.py
@@ -309,3 +309,118 @@ def _socket(*args, **kw):
             await create_connection(addr_info, happy_eyeballs_delay=0.3) == mock_socket
         )
     assert mock_socket.family == socket.AF_INET6
+
+
+@pytest.mark.asyncio
+@patch_socket
+async def test_create_connection_ipv6_and_ipv4_happy_eyeballs_first_ipv6_fails(
+    m_socket: ModuleType,
+) -> None:
+    mock_socket = mock.MagicMock(
+        family=socket.AF_INET,
+        type=socket.SOCK_STREAM,
+        proto=socket.IPPROTO_TCP,
+        fileno=mock.MagicMock(return_value=1),
+    )
+    create_calls = []
+
+    def _socket(*args, **kw):
+        for attr in kw:
+            setattr(mock_socket, attr, kw[attr])
+        return mock_socket
+
+    async def _sock_connect(
+        sock: socket.socket, address: Tuple[str, int, int, int]
+    ) -> None:
+        create_calls.append(address)
+        if address[0] == "dead:beef::":
+            raise OSError("ipv6 fail")
+
+        return None
+
+    m_socket.socket = _socket  # type: ignore
+    ipv6_addr_info = (
+        socket.AF_INET6,
+        socket.SOCK_STREAM,
+        socket.IPPROTO_TCP,
+        "",
+        ("dead:beef::", 80, 0, 0),
+    )
+    ipv6_addr_info_2 = (
+        socket.AF_INET6,
+        socket.SOCK_STREAM,
+        socket.IPPROTO_TCP,
+        "",
+        ("dead:aaaa::", 80, 0, 0),
+    )
+    ipv4_addr_info = (
+        socket.AF_INET,
+        socket.SOCK_STREAM,
+        socket.IPPROTO_TCP,
+        "",
+        ("107.6.106.83", 80),
+    )
+    addr_info = [ipv6_addr_info, ipv6_addr_info_2, ipv4_addr_info]
+    loop = asyncio.get_running_loop()
+    with mock.patch.object(loop, "sock_connect", _sock_connect):
+        assert (
+            await create_connection(addr_info, happy_eyeballs_delay=0.3) == mock_socket
+        )
+
+    # IPv6 addresses are tried first, but the first one fails so IPv4 wins
+    assert mock_socket.family == socket.AF_INET
+    assert create_calls == [("dead:beef::", 80, 0, 0), ("107.6.106.83", 80)]
+
+
+@pytest.mark.asyncio
+@patch_socket
+async def test_create_connection_ipv6_only_happy_eyeballs_first_ipv6_fails(
+    m_socket: ModuleType,
+) -> None:
+    mock_socket = mock.MagicMock(
+        family=socket.AF_INET,
+        type=socket.SOCK_STREAM,
+        proto=socket.IPPROTO_TCP,
+        fileno=mock.MagicMock(return_value=1),
+    )
+    create_calls = []
+
+    def _socket(*args, **kw):
+        for attr in kw:
+            setattr(mock_socket, attr, kw[attr])
+        return mock_socket
+
+    async def _sock_connect(
+        sock: socket.socket, address: Tuple[str, int, int, int]
+    ) -> None:
+        create_calls.append(address)
+        if address[0] == "dead:beef::":
+            raise OSError("ipv6 fail")
+
+        return None
+
+    m_socket.socket = _socket  # type: ignore
+    ipv6_addr_info = (
+        socket.AF_INET6,
+        socket.SOCK_STREAM,
+        socket.IPPROTO_TCP,
+        "",
+        ("dead:beef::", 80, 0, 0),
+    )
+    ipv6_addr_info_2 = (
+        socket.AF_INET6,
+        socket.SOCK_STREAM,
+        socket.IPPROTO_TCP,
+        "",
+        ("dead:aaaa::", 80, 0, 0),
+    )
+    addr_info = [ipv6_addr_info, ipv6_addr_info_2]
+    loop = asyncio.get_running_loop()
+    with mock.patch.object(loop, "sock_connect", _sock_connect):
+        assert (
+            await create_connection(addr_info, happy_eyeballs_delay=0.3) == mock_socket
+        )
+
+    # IPv6 address are tried first, but the first one fails so second IPv6 wins
+    assert mock_socket.family == socket.AF_INET6
+    assert create_calls == [("dead:beef::", 80, 0, 0), ("dead:aaaa::", 80, 0, 0)]