1 files changed, 80 insertions, 0 deletions
diff --git a/tests/test_rdap.py b/tests/test_rdap.py
index b7baa21..9be5bcd 100644
--- a/tests/test_rdap.py
+++ b/tests/test_rdap.py
@@ -261,5 +261,85 @@ class AbuseExtraction(unittest.TestCase):
         )
 
 
+class Query(unittest.TestCase):
+    IPV4 = {"services": [[["198.51.100.0/24"], ["https://rir.example.invalid/"]]]}
+    DNS = {"services": [[["invalid"], ["https://registry.example.invalid/"]]]}
+
+    def test_an_ip_query_hits_the_selected_server(self):
+        calls = []
+
+        def fetch(url):
+            calls.append(url)
+            return {"handle": "NET-1", "entities": []}
+
+        result = rdap.query_ip("198.51.100.7", self.IPV4, fetch=fetch)
+
+        self.assertEqual(calls, ["https://rir.example.invalid/ip/198.51.100.7"])
+        self.assertEqual(result["handle"], "NET-1")
+
+    def test_an_unlisted_ip_is_not_queried(self):
+        def fetch(url):
+            raise AssertionError(f"should not have fetched {url}")
+
+        self.assertIsNone(rdap.query_ip("203.0.113.9", self.IPV4, fetch=fetch))
+
+    def test_the_label_walk_stops_at_the_first_answer(self):
+        """mail.deep.example.invalid is not registrable; example.invalid is.
+
+        The registry is the authority on what is registrable, which is why
+        this walks rather than carrying a Public Suffix List that would go
+        stale weekly.
+        """
+        calls = []
+
+        def fetch(url):
+            calls.append(url)
+            if url.endswith("/domain/example.invalid"):
+                return {"handle": "DOM-1", "entities": []}
+            raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
+
+        result, queried = rdap.query_domain(
+            "mail.deep.example.invalid", self.DNS, fetch=fetch
+        )
+
+        self.assertEqual(queried, "example.invalid")
+        self.assertEqual(result["handle"], "DOM-1")
+        self.assertEqual(len(calls), 3)
+
+    def test_the_walk_never_queries_a_bare_tld(self):
+        calls = []
+
+        def fetch(url):
+            calls.append(url)
+            raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
+
+        result, queried = rdap.query_domain(
+            "deep.example.invalid", self.DNS, fetch=fetch
+        )
+
+        self.assertIsNone(result)
+        self.assertNotIn("https://registry.example.invalid/domain/invalid", calls)
+
+    def test_a_tld_with_no_server_is_not_queried(self):
+        def fetch(url):
+            raise AssertionError(f"should not have fetched {url}")
+
+        result, queried = rdap.query_domain(
+            "example.test", self.DNS, fetch=fetch
+        )
+        self.assertIsNone(result)
+
+    def test_the_walk_is_capped(self):
+        calls = []
+
+        def fetch(url):
+            calls.append(url)
+            raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
+
+        host = "a.b.c.d.e.f.g.example.invalid"
+        rdap.query_domain(host, self.DNS, fetch=fetch)
+        self.assertLessEqual(len(calls), rdap._MAX_LABEL_WALK)
+
+
 if __name__ == "__main__":
     unittest.main()