Harden the DDG collector with token-bucket pacing and prefer the ddgs package

tools/test_refcollect.py

@@ -1,7 +1,12 @@
 """Tests for the local-only reference collector (no network)."""
 
+import datetime
+import email.utils
+import gzip
 import json
+import socket
 import unittest
+import urllib.error
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest import mock
@@ -9,6 +14,37 @@ from unittest import mock
 import refcollect as rc
 
 
+def _fake_clock(start=0.0):
+    """Return (clock, monotonic, sleep) that advance virtual time on sleep."""
+    clock = {"now": start}
+
+    def monotonic():
+        return clock["now"]
+
+    def sleep(seconds):
+        clock["now"] += seconds
+
+    return clock, monotonic, sleep
+
+
+class _FakeResponse:
+    """Minimal stand-in for the urlopen context manager."""
+
+    def __init__(self, body=b"x", content_encoding=None):
+        self._body = body
+        self.headers = ({"Content-Encoding": content_encoding}
+                        if content_encoding else {})
+
+    def read(self):
+        return self._body
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+
 class SlugTests(unittest.TestCase):
     def test_slug_folds_and_strips(self):
         self.assertEqual(rc._slug("El Capitan Entertainment Centre"),
@@ -39,7 +75,8 @@ class SheetTests(unittest.TestCase):
             slug_dir = root / "block"
             slug_dir.mkdir()
             for i in range(3):
-                Image.new("RGB", (640, 400), (i * 60, 0, 0)).save(slug_dir / f"{i:02d}.jpg")
+                Image.new("RGB", (640, 400), (i * 60, 0, 0)
+                          ).save(slug_dir / f"{i:02d}.jpg")
             (slug_dir / "manifest.json").write_text(json.dumps([
                 {"file": "00.jpg"}, {"file": "01.jpg"}, {"file": "02.jpg"}]))
             with mock.patch.object(rc, "RESEARCH_DIR", root):
@@ -47,5 +84,131 @@ class SheetTests(unittest.TestCase):
             self.assertTrue((slug_dir / "sheet.jpg").is_file())
 
 
+class TokenBucketTests(unittest.TestCase):
+    def test_burst_is_spent_before_sleeping(self):
+        clock, monotonic, sleep = _fake_clock()
+        with mock.patch.object(rc.time, "monotonic", monotonic), \
+                mock.patch.object(rc.time, "sleep", sleep):
+            bucket = rc.TokenBucket(rate=1.0, capacity=2.0)
+            self.assertEqual(bucket.acquire(), 0.0)
+            self.assertEqual(bucket.acquire(), 0.0)
+            self.assertAlmostEqual(bucket.acquire(), 1.0)
+            self.assertAlmostEqual(bucket.acquire(), 1.0)
+        self.assertAlmostEqual(clock["now"], 2.0)
+
+    def test_tokens_refill_up_to_capacity_only(self):
+        clock, monotonic, sleep = _fake_clock()
+        with mock.patch.object(rc.time, "monotonic", monotonic), \
+                mock.patch.object(rc.time, "sleep", sleep):
+            bucket = rc.TokenBucket(rate=1.0, capacity=2.0)
+            bucket.acquire(2.0)
+            self.assertAlmostEqual(bucket.tokens, 0.0)
+            clock["now"] += 100.0
+            self.assertAlmostEqual(bucket.tokens, 2.0)  # capped at capacity
+            clock["now"] += 100.0
+            self.assertAlmostEqual(bucket.tokens, 2.0)
+
+    def test_invalid_configuration_is_rejected(self):
+        with self.assertRaises(ValueError):
+            rc.TokenBucket(rate=0)
+        with self.assertRaises(ValueError):
+            rc.TokenBucket(rate=1.0, capacity=0.5)
+
+
+class BackoffTests(unittest.TestCase):
+    def test_backoff_grows_and_stays_within_its_window(self):
+        for attempt in range(1, 6):
+            ceiling = min(rc.BACKOFF_CAP_S, rc.BACKOFF_BASE_S *
+                          2 ** (attempt - 1))
+            for _ in range(20):
+                delay = rc._backoff_delay(attempt)
+                self.assertGreaterEqual(delay, ceiling / 2)
+                self.assertLessEqual(delay, ceiling)
+
+    def test_backoff_is_capped(self):
+        self.assertLessEqual(rc._backoff_delay(20), rc.BACKOFF_CAP_S)
+
+    def test_retry_after_overrides_the_exponential_guess(self):
+        self.assertEqual(rc._backoff_delay(1, retry_after=30.0), 30.0)
+        self.assertEqual(rc._backoff_delay(
+            1, retry_after=0.0), 1.0)  # never instant
+        self.assertEqual(rc._backoff_delay(
+            1, retry_after=10_000.0), rc.BACKOFF_CAP_S)
+
+    def test_retry_after_parses_seconds_and_http_dates(self):
+        self.assertEqual(rc._parse_retry_after("7"), 7.0)
+        self.assertIsNone(rc._parse_retry_after(None))
+        self.assertIsNone(rc._parse_retry_after("soon"))
+        when = email.utils.format_datetime(
+            datetime.datetime.now(datetime.timezone.utc)
+            + datetime.timedelta(seconds=30))
+        self.assertAlmostEqual(rc._parse_retry_after(when), 30.0, delta=2.0)
+
+
+class PolitenessTests(unittest.TestCase):
+    def test_user_agent_is_the_current_firefox_string_on_linux(self):
+        self.assertEqual(rc.BROWSER_HEADERS["User-Agent"], rc.USER_AGENT)
+        self.assertIn("X11; Linux x86_64", rc.USER_AGENT)
+        self.assertIn("Firefox/", rc.USER_AGENT)
+        # A UA on its own is a tell: the supporting browser headers must match.
+        for header in ("Accept", "Accept-Language", "Sec-Fetch-Dest", "Referer"):
+            self.assertIn(header, rc.BROWSER_HEADERS)
+
+    def test_gzip_bodies_are_decoded(self):
+        payload = gzip.compress(b"jpeg-bytes")
+        self.assertEqual(rc._decode_body(payload, "gzip"), b"jpeg-bytes")
+        self.assertEqual(rc._decode_body(b"plain", None), b"plain")
+
+    def test_polite_get_retries_transient_errors_then_succeeds(self):
+        error = urllib.error.HTTPError("https://x/i.jpg", 429,
+                                       "Too Many Requests",
+                                       {"Retry-After": "3"}, None)
+        bucket = rc.TokenBucket(rate=1000, capacity=1000)
+        with mock.patch.object(rc.urllib.request, "urlopen",
+                               side_effect=[error, _FakeResponse(b"jpeg")]) as opened, \
+                mock.patch.object(rc.time, "sleep") as sleep:
+            body, _ = rc.polite_get("https://x/i.jpg", bucket=bucket)
+        self.assertEqual(body, b"jpeg")
+        self.assertEqual(opened.call_count, 2)
+        # Retry-After is honoured instead of the exponential guess.
+        sleep.assert_called_once_with(3.0)
+        sent = opened.call_args_list[0][0][0]
+        self.assertEqual(sent.get_header("User-agent"), rc.USER_AGENT)
+
+    def test_polite_get_does_not_retry_permanent_errors(self):
+        error = urllib.error.HTTPError(
+            "https://x/i.jpg", 404, "Not Found", {}, None)
+        bucket = rc.TokenBucket(rate=1000, capacity=1000)
+        with mock.patch.object(rc.urllib.request, "urlopen",
+                               side_effect=error) as opened, \
+                mock.patch.object(rc.time, "sleep") as sleep, \
+                self.assertRaises(urllib.error.HTTPError):
+            rc.polite_get("https://x/i.jpg", bucket=bucket)
+        self.assertEqual(opened.call_count, 1)
+        sleep.assert_not_called()
+
+    def test_polite_get_does_not_retry_dns_failures(self):
+        error = urllib.error.URLError(socket.gaierror(-2, "Name not known"))
+        bucket = rc.TokenBucket(rate=1000, capacity=1000)
+        with mock.patch.object(rc.urllib.request, "urlopen",
+                               side_effect=error) as opened, \
+                mock.patch.object(rc.time, "sleep") as sleep, \
+                self.assertRaises(urllib.error.URLError):
+            rc.polite_get("https://nope.invalid/i.jpg", bucket=bucket)
+        self.assertEqual(opened.call_count, 1)
+        sleep.assert_not_called()
+
+    def test_polite_get_gives_up_after_the_attempt_budget(self):
+        error = urllib.error.HTTPError(
+            "https://x/i.jpg", 503, "Unavailable", {}, None)
+        bucket = rc.TokenBucket(rate=1000, capacity=1000)
+        with mock.patch.object(rc.urllib.request, "urlopen",
+                               side_effect=error) as opened, \
+                mock.patch.object(rc.time, "sleep"), \
+                self.assertRaises(urllib.error.HTTPError):
+            rc.polite_get("https://x/i.jpg", bucket=bucket, attempts=3)
+        self.assertEqual(opened.call_count, 3)
+
+
 if __name__ == "__main__":
     unittest.main()