ref: make override_blocklist a contextmanager + stackable decorator (#63139)

popped from another branch where I'm adding a `@pytest.mark` above it
which did not play nicely -- opted to make the language here more
inclusive as well

<!-- Describe your PR here. -->

src/sentry/testutils/helpers/socket.py

@@ -1,25 +1,25 @@
+from __future__ import annotations
+
+import contextlib
 import ipaddress
+from typing import Generator
+from unittest import mock
 
 from sentry.net import socket as net_socket
 
-__all__ = ["override_blacklist"]
-
-
-def override_blacklist(*ip_addresses):
-    def decorator(func):
-        def wrapper(*args, **kwargs):
-            disallowed_ips = frozenset(net_socket.DISALLOWED_IPS)
-            net_socket.DISALLOWED_IPS = frozenset(
-                ipaddress.ip_network(str(ip)) for ip in ip_addresses
-            )
-            try:
-                func(*args, **kwargs)
-            finally:
-                net_socket.DISALLOWED_IPS = disallowed_ips
-                # We end up caching these disallowed ips on this function, so
-                # make sure we clear the cache as part of cleanup
-                net_socket.is_ipaddress_allowed.cache_clear()
+__all__ = ["override_blocklist"]
 
-        return wrapper
 
-    return decorator
+@contextlib.contextmanager
+def override_blocklist(*ip_addresses: str) -> Generator[None, None, None]:
+    with mock.patch.object(
+        net_socket,
+        "DISALLOWED_IPS",
+        frozenset(ipaddress.ip_network(ip) for ip in ip_addresses),
+    ):
+        try:
+            yield
+        finally:
+            # We end up caching these disallowed ips on this function, so
+            # make sure we clear the cache as part of cleanup
+            net_socket.is_ipaddress_allowed.cache_clear()

tests/sentry/net/test_socket.py

@@ -4,11 +4,11 @@ from django.test import override_settings
 
 from sentry.net.socket import ensure_fqdn, is_ipaddress_allowed, is_safe_hostname
 from sentry.testutils.cases import TestCase
-from sentry.testutils.helpers import override_blacklist
+from sentry.testutils.helpers import override_blocklist
 
 
 class SocketTest(TestCase):
-    @override_blacklist("10.0.0.0/8", "127.0.0.1")
+    @override_blocklist("10.0.0.0/8", "127.0.0.1")
     def test_is_ipaddress_allowed(self):
         is_ipaddress_allowed.cache_clear()
         assert is_ipaddress_allowed("127.0.0.1") is False
@@ -17,7 +17,7 @@ class SocketTest(TestCase):
         is_ipaddress_allowed.cache_clear()
         assert is_ipaddress_allowed("1.1.1.1") is True
 
-    @override_blacklist("::ffff:10.0.0.0/104", "::1/128")
+    @override_blocklist("::ffff:10.0.0.0/104", "::1/128")
     def test_is_ipaddress_allowed_ipv6(self):
         is_ipaddress_allowed.cache_clear()
         assert is_ipaddress_allowed("::1") is False
@@ -28,7 +28,7 @@ class SocketTest(TestCase):
         is_ipaddress_allowed.cache_clear()
         assert is_ipaddress_allowed("2001:db8:a::123") is True
 
-    @override_blacklist("10.0.0.0/8", "127.0.0.1")
+    @override_blocklist("10.0.0.0/8", "127.0.0.1")
     @patch("socket.getaddrinfo")
     def test_is_safe_hostname(self, mock_getaddrinfo):
         mock_getaddrinfo.return_value = [(2, 1, 6, "", ("81.0.0.1", 0))]

tests/sentry/shared_integrations/client/test_base.py

@@ -9,7 +9,7 @@ from sentry.net.http import Session
 from sentry.shared_integrations.client.base import BaseApiClient
 from sentry.shared_integrations.exceptions import ApiHostError
 from sentry.testutils.cases import TestCase
-from sentry.testutils.helpers.socket import override_blacklist
+from sentry.testutils.helpers.socket import override_blocklist
 
 
 class BaseApiClientTest(TestCase):
@@ -66,7 +66,7 @@ class BaseApiClientTest(TestCase):
     @responses.activate
     @patch.object(BaseApiClient, "finalize_request", side_effect=lambda req: req)
     @patch.object(Session, "send", side_effect=RestrictedIPAddress())
-    @override_blacklist("172.16.0.0/12")
+    @override_blocklist("172.16.0.0/12")
     def test_restricted_ip_address(self, mock_finalize_request, mock_session_send):
         assert not mock_finalize_request.called
         with raises(ApiHostError):

tests/sentry/test_http.py

@@ -9,7 +9,7 @@ from django.core.exceptions import SuspiciousOperation
 from urllib3.util.connection import HAS_IPV6
 
 from sentry import http
-from sentry.testutils.helpers import override_blacklist
+from sentry.testutils.helpers import override_blocklist
 
 
 @responses.activate
@@ -27,7 +27,7 @@ def test_simple(mock_getaddrinfo):
     assert "gzip" in request.headers.get("Accept-Encoding", "")
 
 
-@override_blacklist("127.0.0.1", "::1", "10.0.0.0/8")
+@override_blocklist("127.0.0.1", "::1", "10.0.0.0/8")
 # XXX(dcramer): we can't use responses here as it hooks Session.send
 # @responses.activate
 def test_ip_blacklist_ipv4():
@@ -41,14 +41,14 @@ def test_ip_blacklist_ipv4():
 
 
 @pytest.mark.skipif(not HAS_IPV6, reason="needs ipv6")
-@override_blacklist("::1")
+@override_blocklist("::1")
 def test_ip_blacklist_ipv6():
     with pytest.raises(SuspiciousOperation):
         http.safe_urlopen("http://[::1]")
 
 
 @pytest.mark.skipif(HAS_IPV6, reason="stub for non-ipv6 systems")
-@override_blacklist("::1")
+@override_blocklist("::1")
 @patch("socket.getaddrinfo")
 def test_ip_blacklist_ipv6_fallback(mock_getaddrinfo):
     mock_getaddrinfo.return_value = [(10, 1, 6, "", ("::1", 0, 0, 0))]
@@ -59,7 +59,7 @@ def test_ip_blacklist_ipv6_fallback(mock_getaddrinfo):
 @pytest.mark.skipif(
     platform.system() == "Darwin", reason="macOS is always broken, see comment in sentry/http.py"
 )
-@override_blacklist("127.0.0.1")
+@override_blocklist("127.0.0.1")
 def test_garbage_ip():
     with pytest.raises(SuspiciousOperation):
         # '0177.0000.0000.0001' is an octal for '127.0.0.1'