1 files changed, 5 insertions, 13 deletions
diff --git a/tests/test_url_security.py b/tests/test_url_security.py
index fba6d7d..b10e6ff 100644
--- a/tests/test_url_security.py
+++ b/tests/test_url_security.py
@@ -1,3 +1,5 @@
+from unittest.mock import Mock, patch
+
 import pytest
 
 from webmentions_ssg.url_security import (
@@ -6,14 +8,7 @@ from webmentions_ssg.url_security import (
 )
 
 
-class FakeAnswers:
-    def __init__(self, *addresses: str):
-        self._addresses = addresses
-
-    def addresses(self):
-        return iter(self._addresses)
-
-
+@patch("webmentions_ssg.url_security.dns.resolver.resolve_name")
 @pytest.mark.parametrize(
     ("url", "resolved_addresses"),
     [
@@ -29,14 +24,11 @@ class FakeAnswers:
     ],
 )
 def test_ensure_public_url_rejects_non_public_addresses(
-    monkeypatch: pytest.MonkeyPatch,
+    resolve_name: Mock,
     url: str,
     resolved_addresses: list[str],
 ) -> None:
-    monkeypatch.setattr(
-        "webmentions_ssg.url_security.dns.resolver.resolve_name",
-        lambda hostname: FakeAnswers(*resolved_addresses),
-    )
+    resolve_name.return_value.addresses.return_value = resolved_addresses
 
     with pytest.raises(NonPublicAddressError):
         ensure_public_url(url)