In tests replace live DNS queries with mocked DNS answers that have been captured to a JSON file

JoshData · JoshData · commit d40f70881e62 · 2023-02-04T18:24:49.000-05:00
diff --git a/README.md b/README.md
@@ -416,6 +416,8 @@ pip install -r test_requirements.txt
 make test
 ```
 
+Tests run with mocked DNS responses. When adding or changing tests, temporarily turn on the `BUILD_MOCKED_DNS_RESPONSE_DATA` flag in `tests/mocked_dns_responses.py` to re-build the database of mocked responses from live queries.
+
 For Project Maintainers
 -----------------------
 
diff --git a/email_validator/__main__.py b/email_validator/__main__.py
@@ -17,10 +17,12 @@
 from .exceptions_types import EmailNotValidError
 
 
-def main():
+def main(dns_resolver=None):
+    # The dns_resolver argument is for tests.
+
     if len(sys.argv) == 1:
         # Validate the email addresses pased line-by-line on STDIN.
-        dns_resolver = caching_resolver()
+        dns_resolver = dns_resolver or caching_resolver()
         for line in sys.stdin:
             email = line.strip()
             try:
@@ -31,7 +33,7 @@ def main():
         # Validate the email address passed on the command line.
         email = sys.argv[1]
         try:
-            result = validate_email(email)
+            result = validate_email(email, dns_resolver=dns_resolver)
             print(json.dumps(result.as_dict(), indent=2, sort_keys=True, ensure_ascii=False))
         except EmailNotValidError as e:
             print(e)
diff --git a/email_validator/deliverability.py b/email_validator/deliverability.py
@@ -33,12 +33,6 @@ def validate_email_deliverability(domain, domain_i18n, timeout=None, dns_resolve
     deliverability_info = {}
 
     try:
-        # We need a way to check how timeouts are handled in the tests. So we
-        # have a secret variable that if set makes this method always test the
-        # handling of a timeout.
-        if getattr(validate_email_deliverability, 'TEST_CHECK_TIMEOUT', False):
-            raise dns.exception.Timeout()
-
         try:
             # Try resolving for MX records.
             response = dns_resolver.resolve(domain, "MX")
diff --git a/tests/mocked-dns-answers.json b/tests/mocked-dns-answers.json
@@ -0,0 +1,100 @@
+[
+ {
+  "query": {
+   "name": "gmail.com",
+   "type": "MX",
+   "class": "IN"
+  },
+  "answer": [
+   "10 alt1.gmail-smtp-in.l.google.com.",
+   "30 alt3.gmail-smtp-in.l.google.com.",
+   "5 gmail-smtp-in.l.google.com.",
+   "20 alt2.gmail-smtp-in.l.google.com.",
+   "40 alt4.gmail-smtp-in.l.google.com."
+  ]
+ },
+ {
+  "query": {
+   "name": "xkxufoekjvjfjeodlfmdfjcu.com",
+   "type": "ANY",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "xkxufoekjvjfjeodlfmdfjcu.com",
+   "type": "AAAA",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "example.com",
+   "type": "MX",
+   "class": "IN"
+  },
+  "answer": [
+   "0 ."
+  ]
+ },
+ {
+  "query": {
+   "name": "mail.example",
+   "type": "MX",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "mail.example",
+   "type": "ANY",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "mail.example",
+   "type": "AAAA",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "mail.example.com",
+   "type": "MX",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "mail.example.com",
+   "type": "ANY",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "mail.example.com",
+   "type": "AAAA",
+   "class": "IN"
+  },
+  "answer": []
+ },
+ {
+  "query": {
+   "name": "google.com",
+   "type": "MX",
+   "class": "IN"
+  },
+  "answer": [
+   "10 smtp.google.com."
+  ]
+ }
+]
diff --git a/tests/mocked_dns_response.py b/tests/mocked_dns_response.py
@@ -0,0 +1,120 @@
+import dns.resolver
+import json
+import os.path
+import pytest
+
+from email_validator.deliverability import caching_resolver
+
+# To run deliverability checks without actually making
+# DNS queries, we use a caching resolver where the cache
+# is pre-loaded with DNS responses.
+
+# When False, all DNS queries must come from the mocked
+# data. When True, tests are run with live DNS queries
+# and the DNS responses are saved to a file.
+BUILD_MOCKED_DNS_RESPONSE_DATA = False
+
+
+# This class implements the 'get' and 'put' methods
+# expected for a dns.resolver.Resolver's cache.
+class MockedDnsResponseData:
+    DATA_PATH = os.path.dirname(__file__) + "/mocked-dns-answers.json"
+
+    @staticmethod
+    def create_resolver():
+        if not hasattr(MockedDnsResponseData, 'INSTANCE'):
+            # Create a singleton instance of this class and load the saved DNS responses.
+            # Except when BUILD_MOCKED_DNS_RESPONSE_DATA is true, don't load the data.
+            singleton = MockedDnsResponseData()
+            if not BUILD_MOCKED_DNS_RESPONSE_DATA:
+                singleton.load()
+            MockedDnsResponseData.INSTANCE = singleton
+
+        # Return a new dns.resolver.Resolver configured for caching
+        # using the singleton instance.
+        return caching_resolver(cache=MockedDnsResponseData.INSTANCE)
+
+    def __init__(self):
+        self.data = {}
+
+    def load(self):
+        # Loads the saved DNS response data from the JSON file and
+        # re-structures it into dnspython classes.
+        class Ans:  # mocks the dns.resolver.Answer class
+
+            def __init__(self, rrset):
+                self.rrset = rrset
+
+            def __iter__(self):
+                return iter(self.rrset)
+
+        with open(self.DATA_PATH) as f:
+            data = json.load(f)
+            for item in data:
+                key = (dns.name.from_text(item["query"]["name"] + "."),
+                       dns.rdatatype.from_text(item["query"]["type"]),
+                       dns.rdataclass.from_text(item["query"]["class"]))
+                rdatas = [
+                    dns.rdata.from_text(rdtype=key[1], rdclass=key[2], tok=rr)
+                    for rr in item["answer"]
+                ]
+                if item["answer"]:
+                    self.data[key] = Ans(dns.rdataset.from_rdata_list(0, rdatas=rdatas))
+                else:
+                    self.data[key] = None
+
+    def save(self):
+        # Re-structure as a list with basic data types.
+        data = [
+            {
+                "query": {
+                    "name": key[0].to_text(omit_final_dot=True),
+                    "type": dns.rdatatype.to_text(key[1]),
+                    "class": dns.rdataclass.to_text(key[2]),
+                },
+                "answer": [
+                    rr.to_text()
+                    for rr in value
+                ]
+            }
+            for key, value in self.data.items()
+        ]
+        with open(self.DATA_PATH, "w") as f:
+            json.dump(data, f, indent=True)
+
+    def get(self, key):
+        # Special-case a domain to create a timeout.
+        if key[0].to_text() == "timeout.com.":
+            raise dns.exception.Timeout()
+
+        # When building the DNS response database, return
+        # a cache miss.
+        if BUILD_MOCKED_DNS_RESPONSE_DATA:
+            return None
+
+        # Query the data for a matching record.
+        if key in self.data:
+            if not self.data[key]:
+                raise dns.resolver.NoAnswer()
+            return self.data[key]
+
+        # Query the data for a response to an ANY query.
+        ANY = dns.rdatatype.from_text("ANY")
+        if (key[0], ANY, key[2]) in self.data and self.data[(key[0], ANY, key[2])] is None:
+            raise dns.resolver.NoAnswer()
+
+        raise ValueError("Saved DNS data did not contain query: {}".format(key))
+
+    def put(self, key, value):
+        # Build the DNS data by saving the live query response.
+        if not BUILD_MOCKED_DNS_RESPONSE_DATA:
+            raise ValueError("Should not get here.")
+        self.data[key] = value
+
+
+@pytest.fixture(scope="session", autouse=True)
+def MockedDnsResponseDataCleanup(request):
+    def cleanup_func():
+        if BUILD_MOCKED_DNS_RESPONSE_DATA:
+            MockedDnsResponseData.INSTANCE.save()
+    request.addfinalizer(cleanup_func)
diff --git a/tests/test_deliverability.py b/tests/test_deliverability.py
@@ -1,14 +1,17 @@
-import dns.resolver
 import pytest
 import re
 
 from email_validator import EmailUndeliverableError, \
                             validate_email
-from email_validator.deliverability import caching_resolver, validate_email_deliverability
+from email_validator.deliverability import validate_email_deliverability
+
+from mocked_dns_response import MockedDnsResponseData, MockedDnsResponseDataCleanup  # noqa: F401
+
+RESOLVER = MockedDnsResponseData.create_resolver()
 
 
 def test_deliverability_found():
-    response = validate_email_deliverability('gmail.com', 'gmail.com')
+    response = validate_email_deliverability('gmail.com', 'gmail.com', dns_resolver=RESOLVER)
     assert response.keys() == {'mx', 'mx_fallback_type'}
     assert response['mx_fallback_type'] is None
     assert len(response['mx']) > 1
@@ -21,12 +24,12 @@ def test_deliverability_fails():
     # No MX record.
     domain = 'xkxufoekjvjfjeodlfmdfjcu.com'
     with pytest.raises(EmailUndeliverableError, match='The domain name {} does not exist'.format(domain)):
-        validate_email_deliverability(domain, domain)
+        validate_email_deliverability(domain, domain, dns_resolver=RESOLVER)
 
     # Null MX record.
     domain = 'example.com'
     with pytest.raises(EmailUndeliverableError, match='The domain name {} does not accept email'.format(domain)):
-        validate_email_deliverability(domain, domain)
+        validate_email_deliverability(domain, domain, dns_resolver=RESOLVER)
 
 
 @pytest.mark.parametrize(
@@ -41,48 +44,12 @@ def test_email_example_reserved_domain(email_input):
     # Since these all fail deliverabiltiy from a static list,
     # DNS deliverability checks do not arise.
     with pytest.raises(EmailUndeliverableError) as exc_info:
-        validate_email(email_input)
+        validate_email(email_input, dns_resolver=RESOLVER)
     # print(f'({email_input!r}, {str(exc_info.value)!r}),')
     assert re.match(r"The domain name [a-z\.]+ does not (accept email|exist)\.", str(exc_info.value)) is not None
 
 
 def test_deliverability_dns_timeout():
-    validate_email_deliverability.TEST_CHECK_TIMEOUT = True
-    response = validate_email_deliverability('gmail.com', 'gmail.com')
+    response = validate_email_deliverability('timeout.com', 'timeout.com', dns_resolver=RESOLVER)
     assert "mx" not in response
     assert response.get("unknown-deliverability") == "timeout"
-    validate_email('test@gmail.com')
-    del validate_email_deliverability.TEST_CHECK_TIMEOUT
-
-
-def test_validate_email__with_caching_resolver():
-    # unittest.mock.patch("dns.resolver.LRUCache.get") doesn't
-    # work --- it causes get to always return an empty list.
-    # So we'll mock our own way.
-    class MockedCache:
-        get_called = False
-        put_called = False
-
-        def get(self, key):
-            self.get_called = True
-            return None
-
-        def put(self, key, value):
-            self.put_called = True
-
-    # Test with caching_resolver helper method.
-    mocked_cache = MockedCache()
-    dns_resolver = caching_resolver(cache=mocked_cache)
-    validate_email("test@gmail.com", dns_resolver=dns_resolver)
-    assert mocked_cache.put_called
-    validate_email("test@gmail.com", dns_resolver=dns_resolver)
-    assert mocked_cache.get_called
-
-    # Test with dns.resolver.Resolver instance.
-    dns_resolver = dns.resolver.Resolver()
-    dns_resolver.lifetime = 10
-    dns_resolver.cache = MockedCache()
-    validate_email("test@gmail.com", dns_resolver=dns_resolver)
-    assert mocked_cache.put_called
-    validate_email("test@gmail.com", dns_resolver=dns_resolver)
-    assert mocked_cache.get_called
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -2,6 +2,10 @@
 # Let's test main but rename it to be clear
 from email_validator.__main__ import main as validator_command_line_tool
 
+from mocked_dns_response import MockedDnsResponseData, MockedDnsResponseDataCleanup  # noqa: F401
+
+RESOLVER = MockedDnsResponseData.create_resolver()
+
 
 def test_dict_accessor():
     input_email = "testaddr@example.tld"
@@ -14,17 +18,17 @@ def test_main_single_good_input(monkeypatch, capsys):
     import json
     test_email = "google@google.com"
     monkeypatch.setattr('sys.argv', ['email_validator', test_email])
-    validator_command_line_tool()
+    validator_command_line_tool(dns_resolver=RESOLVER)
     stdout, _ = capsys.readouterr()
     output = json.loads(str(stdout))
     assert isinstance(output, dict)
-    assert validate_email(test_email).original_email == output["original_email"]
+    assert validate_email(test_email, dns_resolver=RESOLVER).original_email == output["original_email"]
 
 
 def test_main_single_bad_input(monkeypatch, capsys):
     bad_email = 'test@..com'
     monkeypatch.setattr('sys.argv', ['email_validator', bad_email])
-    validator_command_line_tool()
+    validator_command_line_tool(dns_resolver=RESOLVER)
     stdout, _ = capsys.readouterr()
     assert stdout == 'An email address cannot have a period immediately after the @-sign.\n'
 
@@ -35,7 +39,7 @@ def test_main_multi_input(monkeypatch, capsys):
     test_input = io.StringIO("\n".join(test_cases))
     monkeypatch.setattr('sys.stdin', test_input)
     monkeypatch.setattr('sys.argv', ['email_validator'])
-    validator_command_line_tool()
+    validator_command_line_tool(dns_resolver=RESOLVER)
     stdout, _ = capsys.readouterr()
     assert test_cases[0] not in stdout
     assert test_cases[1] not in stdout
