Break apart new_socket to be testable (#867)

bdraco · web-flow · commit 22ff6b56d7b6 · 2021-06-27T21:47:17.000-10:00
diff --git a/tests/utils/test_net.py b/tests/utils/test_net.py
@@ -8,6 +8,7 @@
 import errno
 import ifaddr
 import pytest
+import socket
 import unittest
 
 from zeroconf._utils import net as netutils
@@ -99,3 +100,42 @@ def test_autodetect_ip_version():
     assert r.autodetect_ip_version([]) is r.IPVersion.V4Only
     assert r.autodetect_ip_version(["::1", "1.2.3.4"]) is r.IPVersion.All
     assert r.autodetect_ip_version(["::1"]) is r.IPVersion.V6Only
+
+
+def test_disable_ipv6_only_or_raise():
+    """Test that IPV6_V6ONLY failing logs a nice error message and still raises."""
+    errors_logged = []
+
+    def _log_error(*args):
+        nonlocal errors_logged
+        errors_logged.append(args)
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    with pytest.raises(OSError), patch.object(netutils.log, "error", _log_error), patch(
+        "socket.socket.setsockopt", side_effect=OSError
+    ):
+        netutils.disable_ipv6_only_or_raise(sock)
+
+    assert (
+        errors_logged[0][0]
+        == 'Support for dual V4-V6 sockets is not present, use IPVersion.V4 or IPVersion.V6'
+    )
+
+
+@pytest.mark.skipif(not hasattr(socket, 'SO_REUSEPORT'), reason="System does not have SO_REUSEPORT")
+def test_set_so_reuseport_if_available_is_present():
+    """Test that setting socket.SO_REUSEPORT only OSError errno.ENOPROTOOPT is trapped."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    with pytest.raises(OSError), patch("socket.socket.setsockopt", side_effect=OSError):
+        netutils.set_so_reuseport_if_available(sock)
+
+    with patch("socket.socket.setsockopt", side_effect=OSError(errno.ENOPROTOOPT, None)):
+        netutils.set_so_reuseport_if_available(sock)
+
+
+@pytest.mark.skipif(hasattr(socket, 'SO_REUSEPORT'), reason="System has SO_REUSEPORT")
+def test_set_so_reuseport_if_available_not_present():
+    """Test that we do not try to set SO_REUSEPORT if it is not present."""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    with patch("socket.socket.setsockopt", side_effect=OSError):
+        netutils.set_so_reuseport_if_available(sock)
diff --git a/zeroconf/_utils/net.py b/zeroconf/_utils/net.py
@@ -161,6 +161,34 @@ def normalize_interface_choice(
     return result
 
 
+def disable_ipv6_only_or_raise(s: socket.socket) -> None:
+    """Make V6 sockets work for both V4 and V6 (required for Windows)."""
+    try:
+        s.setsockopt(_IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
+    except OSError:
+        log.error('Support for dual V4-V6 sockets is not present, use IPVersion.V4 or IPVersion.V6')
+        raise
+
+
+def set_so_reuseport_if_available(s: socket.socket) -> None:
+    """Set SO_REUSEADDR on a socket if available."""
+    # SO_REUSEADDR should be equivalent to SO_REUSEPORT for
+    # multicast UDP sockets (p 731, "TCP/IP Illustrated,
+    # Volume 2"), but some BSD-derived systems require
+    # SO_REUSEPORT to be specified explicitly.  Also, not all
+    # versions of Python have SO_REUSEPORT available.
+    # Catch OSError and socket.error for kernel versions <3.9 because lacking
+    # SO_REUSEPORT support.
+    if not hasattr(socket, 'SO_REUSEPORT'):
+        return
+
+    try:
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)  # pylint: disable=no-member
+    except OSError as err:
+        if err.errno != errno.ENOPROTOOPT:
+            raise
+
+
 def new_socket(  # pylint: disable=too-many-branches
     bind_addr: Union[Tuple[str], Tuple[str, int, int]],
     port: int = _MDNS_PORT,
@@ -180,28 +208,11 @@ def new_socket(  # pylint: disable=too-many-branches
         s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
 
     if ip_version == IPVersion.All:
-        # make V6 sockets work for both V4 and V6 (required for Windows)
-        try:
-            s.setsockopt(_IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
-        except OSError:
-            log.error('Support for dual V4-V6 sockets is not present, use IPVersion.V4 or IPVersion.V6')
-            raise
+        disable_ipv6_only_or_raise(s)
 
     s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 
-    # SO_REUSEADDR should be equivalent to SO_REUSEPORT for
-    # multicast UDP sockets (p 731, "TCP/IP Illustrated,
-    # Volume 2"), but some BSD-derived systems require
-    # SO_REUSEPORT to be specified explicitly.  Also, not all
-    # versions of Python have SO_REUSEPORT available.
-    # Catch OSError and socket.error for kernel versions <3.9 because lacking
-    # SO_REUSEPORT support.
-    if hasattr(socket, 'SO_REUSEPORT'):
-        try:
-            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)  # pylint: disable=no-member
-        except OSError as err:
-            if err.errno != errno.ENOPROTOOPT:
-                raise
+    set_so_reuseport_if_available(s)
 
     if port == _MDNS_PORT:
         ttl = struct.pack(b'B', 255)
