refactored the duplicate mock DNS resolution code into a reusable fixture

Author: acul71

tests/test_resolvers.py

@@ -14,6 +14,7 @@
 if sys.version_info >= (3, 11):
     from builtins import BaseExceptionGroup
 else:
+
     class BaseExceptionGroup(Exception):
         pass
 
@@ -24,17 +25,9 @@ def dns_resolver():
     return DNSResolver()
 
 
-@pytest.mark.trio
-async def test_resolve_non_dns_addr(dns_resolver):
-    """Test resolving a non-DNS multiaddr."""
-    ma = Multiaddr("/ip4/127.0.0.1/tcp/1234")
-    result = await dns_resolver.resolve(ma)
-    assert result == [ma]
-
-
-@pytest.mark.trio
-async def test_resolve_dns_addr(dns_resolver):
-    """Test resolving a DNS multiaddr."""
+@pytest.fixture
+def mock_dns_resolution():
+    """Create mock DNS resolution setup for testing."""
     # Create mock DNS answer for A record (IPv4)
     mock_answer_a = AsyncMock()
     mock_rdata_a = AsyncMock()
@@ -45,17 +38,35 @@ async def test_resolve_dns_addr(dns_resolver):
     mock_answer_aaaa = AsyncMock()
     mock_answer_aaaa.__iter__.return_value = []
 
-    with patch.object(dns_resolver._resolver, 'resolve') as mock_resolve:
-        # Configure the mock to return different results based on record type
-        async def mock_resolve_side_effect(hostname, record_type):
-            if record_type == "A":
-                return mock_answer_a
-            elif record_type == "AAAA":
-                return mock_answer_aaaa
-            else:
-                raise dns.resolver.NXDOMAIN()
+    # Configure the mock to return different results based on record type
+    async def mock_resolve_side_effect(hostname, record_type):
+        if record_type == "A":
+            return mock_answer_a
+        elif record_type == "AAAA":
+            return mock_answer_aaaa
+        else:
+            raise dns.resolver.NXDOMAIN()
+
+    return {
+        "mock_answer_a": mock_answer_a,
+        "mock_answer_aaaa": mock_answer_aaaa,
+        "mock_resolve_side_effect": mock_resolve_side_effect,
+    }
+
+
+@pytest.mark.trio
+async def test_resolve_non_dns_addr(dns_resolver):
+    """Test resolving a non-DNS multiaddr."""
+    ma = Multiaddr("/ip4/127.0.0.1/tcp/1234")
+    result = await dns_resolver.resolve(ma)
+    assert result == [ma]
+
 
-        mock_resolve.side_effect = mock_resolve_side_effect
+@pytest.mark.trio
+async def test_resolve_dns_addr(dns_resolver, mock_dns_resolution):
+    """Test resolving a DNS multiaddr."""
+    with patch.object(dns_resolver._resolver, "resolve") as mock_resolve:
+        mock_resolve.side_effect = mock_dns_resolution["mock_resolve_side_effect"]
 
         ma = Multiaddr("/dnsaddr/example.com")
         result = await dns_resolver.resolve(ma)
@@ -65,29 +76,10 @@ async def mock_resolve_side_effect(hostname, record_type):
 
 
 @pytest.mark.trio
-async def test_resolve_dns_addr_with_peer_id(dns_resolver):
+async def test_resolve_dns_addr_with_peer_id(dns_resolver, mock_dns_resolution):
     """Test resolving a DNS multiaddr with a peer ID."""
-    # Create mock DNS answer for A record (IPv4)
-    mock_answer_a = AsyncMock()
-    mock_rdata_a = AsyncMock()
-    mock_rdata_a.address = "127.0.0.1"
-    mock_answer_a.__iter__.return_value = [mock_rdata_a]
-
-    # Create mock DNS answer for AAAA record (IPv6) - return empty to avoid conflicts
-    mock_answer_aaaa = AsyncMock()
-    mock_answer_aaaa.__iter__.return_value = []
-
-    with patch.object(dns_resolver._resolver, 'resolve') as mock_resolve:
-        # Configure the mock to return different results based on record type
-        async def mock_resolve_side_effect(hostname, record_type):
-            if record_type == "A":
-                return mock_answer_a
-            elif record_type == "AAAA":
-                return mock_answer_aaaa
-            else:
-                raise dns.resolver.NXDOMAIN()
-
-        mock_resolve.side_effect = mock_resolve_side_effect
+    with patch.object(dns_resolver._resolver, "resolve") as mock_resolve:
+        mock_resolve.side_effect = mock_dns_resolution["mock_resolve_side_effect"]
 
         ma = Multiaddr("/dnsaddr/example.com/p2p/QmYyQSo1c1Ym7orWxLYvCrM2EmxFTANf8wXmmE7wjh53Qk")
         result = await dns_resolver.resolve(ma)
@@ -98,29 +90,10 @@ async def mock_resolve_side_effect(hostname, record_type):
 
 
 @pytest.mark.trio
-async def test_resolve_recursive_dns_addr(dns_resolver):
+async def test_resolve_recursive_dns_addr(dns_resolver, mock_dns_resolution):
     """Test resolving a recursive DNS multiaddr."""
-    # Create mock DNS answer for A record (IPv4)
-    mock_answer_a = AsyncMock()
-    mock_rdata_a = AsyncMock()
-    mock_rdata_a.address = "127.0.0.1"
-    mock_answer_a.__iter__.return_value = [mock_rdata_a]
-
-    # Create mock DNS answer for AAAA record (IPv6) - return empty to avoid conflicts
-    mock_answer_aaaa = AsyncMock()
-    mock_answer_aaaa.__iter__.return_value = []
-
-    with patch.object(dns_resolver._resolver, 'resolve') as mock_resolve:
-        # Configure the mock to return different results based on record type
-        async def mock_resolve_side_effect(hostname, record_type):
-            if record_type == "A":
-                return mock_answer_a
-            elif record_type == "AAAA":
-                return mock_answer_aaaa
-            else:
-                raise dns.resolver.NXDOMAIN()
-
-        mock_resolve.side_effect = mock_resolve_side_effect
+    with patch.object(dns_resolver._resolver, "resolve") as mock_resolve:
+        mock_resolve.side_effect = mock_dns_resolution["mock_resolve_side_effect"]
 
         ma = Multiaddr("/dnsaddr/example.com")
         result = await dns_resolver.resolve(ma, {"max_recursive_depth": 2})
@@ -140,37 +113,18 @@ async def test_resolve_recursion_limit(dns_resolver):
 @pytest.mark.trio
 async def test_resolve_dns_addr_error(dns_resolver):
     """Test handling DNS resolution errors."""
-    with patch.object(dns_resolver._resolver, 'resolve', side_effect=dns.resolver.NXDOMAIN):
+    with patch.object(dns_resolver._resolver, "resolve", side_effect=dns.resolver.NXDOMAIN):
         ma = Multiaddr("/dnsaddr/example.com")
         # When DNS resolution fails, the resolver should return the original multiaddr
         result = await dns_resolver.resolve(ma)
         assert result == [ma]
 
 
 @pytest.mark.trio
-async def test_resolve_dns_addr_with_quotes(dns_resolver):
+async def test_resolve_dns_addr_with_quotes(dns_resolver, mock_dns_resolution):
     """Test resolving DNS records with quoted strings."""
-    # Create mock DNS answer for A record (IPv4)
-    mock_answer_a = AsyncMock()
-    mock_rdata_a = AsyncMock()
-    mock_rdata_a.address = "127.0.0.1"
-    mock_answer_a.__iter__.return_value = [mock_rdata_a]
-
-    # Create mock DNS answer for AAAA record (IPv6) - return empty to avoid conflicts
-    mock_answer_aaaa = AsyncMock()
-    mock_answer_aaaa.__iter__.return_value = []
-
-    with patch.object(dns_resolver._resolver, 'resolve') as mock_resolve:
-        # Configure the mock to return different results based on record type
-        async def mock_resolve_side_effect(hostname, record_type):
-            if record_type == "A":
-                return mock_answer_a
-            elif record_type == "AAAA":
-                return mock_answer_aaaa
-            else:
-                raise dns.resolver.NXDOMAIN()
-
-        mock_resolve.side_effect = mock_resolve_side_effect
+    with patch.object(dns_resolver._resolver, "resolve") as mock_resolve:
+        mock_resolve.side_effect = mock_dns_resolution["mock_resolve_side_effect"]
 
         ma = Multiaddr("/dnsaddr/example.com")
         result = await dns_resolver.resolve(ma)
@@ -180,29 +134,10 @@ async def mock_resolve_side_effect(hostname, record_type):
 
 
 @pytest.mark.trio
-async def test_resolve_dns_addr_with_mixed_quotes(dns_resolver):
+async def test_resolve_dns_addr_with_mixed_quotes(dns_resolver, mock_dns_resolution):
     """Test resolving DNS records with mixed quotes."""
-    # Create mock DNS answer for A record (IPv4)
-    mock_answer_a = AsyncMock()
-    mock_rdata_a = AsyncMock()
-    mock_rdata_a.address = "127.0.0.1"
-    mock_answer_a.__iter__.return_value = [mock_rdata_a]
-
-    # Create mock DNS answer for AAAA record (IPv6) - return empty to avoid conflicts
-    mock_answer_aaaa = AsyncMock()
-    mock_answer_aaaa.__iter__.return_value = []
-
-    with patch.object(dns_resolver._resolver, 'resolve') as mock_resolve:
-        # Configure the mock to return different results based on record type
-        async def mock_resolve_side_effect(hostname, record_type):
-            if record_type == "A":
-                return mock_answer_a
-            elif record_type == "AAAA":
-                return mock_answer_aaaa
-            else:
-                raise dns.resolver.NXDOMAIN()
-
-        mock_resolve.side_effect = mock_resolve_side_effect
+    with patch.object(dns_resolver._resolver, "resolve") as mock_resolve:
+        mock_resolve.side_effect = mock_dns_resolution["mock_resolve_side_effect"]
 
         ma = Multiaddr("/dnsaddr/example.com")
         result = await dns_resolver.resolve(ma)
@@ -224,7 +159,7 @@ async def slow_dns_resolve(*args, **kwargs):
         await trio.sleep(0.5)  # Long sleep to allow cancellation
         raise dns.resolver.NXDOMAIN("Domain not found")
 
-    with patch.object(dns_resolver._resolver, 'resolve', side_effect=slow_dns_resolve):
+    with patch.object(dns_resolver._resolver, "resolve", side_effect=slow_dns_resolve):
         # Start resolution in background and cancel it
         async with trio.open_nursery() as nursery:
             # Start the resolution