From 334eb95e568aec20a3914194725d0c7b64f67926 Mon Sep 17 00:00:00 2001
From: Mathieu Leplatre <mathieu@mozilla.com>
Date: Tue, 10 Dec 2024 09:43:59 +0100
Subject: [PATCH] Fix #1447: add option to shard check of attachments (#1517)

---
 .../attachments_availability.py               | 10 +++-
 .../test_attachments_availability.py          | 54 +++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/checks/remotesettings/attachments_availability.py b/checks/remotesettings/attachments_availability.py
index 997b8a45..fd1117ce 100644
--- a/checks/remotesettings/attachments_availability.py
+++ b/checks/remotesettings/attachments_availability.py
@@ -4,6 +4,8 @@
 The URLs of unreachable attachments is returned along with the number of checked records.
 """
 
+import math
+
 import aiohttp
 
 from telescope.typings import CheckResult
@@ -20,7 +22,7 @@ async def test_url(url):
         return False
 
 
-async def run(server: str) -> CheckResult:
+async def run(server: str, slice_percent: tuple[int, int] = (0, 100)) -> CheckResult:
     client = KintoClient(server_url=server)
 
     info = await client.server_info()
@@ -47,7 +49,11 @@ async def run(server: str) -> CheckResult:
                 continue
             url = base_url + record["attachment"]["location"]
             urls.append(url)
-    futures = [test_url(url) for url in urls]
+
+    lower_idx = math.floor(slice_percent[0] / 100.0 * len(urls))
+    upper_idx = math.ceil(slice_percent[1] / 100.0 * len(urls))
+
+    futures = [test_url(url) for url in urls[lower_idx:upper_idx]]
     results = await run_parallel(*futures)
     missing = [url for url, success in zip(urls, results) if not success]
 
diff --git a/tests/checks/remotesettings/test_attachments_availability.py b/tests/checks/remotesettings/test_attachments_availability.py
index 54439763..e484e294 100644
--- a/tests/checks/remotesettings/test_attachments_availability.py
+++ b/tests/checks/remotesettings/test_attachments_availability.py
@@ -1,3 +1,7 @@
+from unittest import mock
+
+import pytest
+
 from checks.remotesettings.attachments_availability import run
 
 
@@ -72,3 +76,53 @@ async def test_negative(mock_responses, mock_aioresponses):
 
     assert status is False
     assert data == {"missing": ["http://cdn/missing.jpg"], "checked": 2}
+
+
+@pytest.mark.parametrize(
+    ("slice_percent", "expected_lower", "expected_upper"),
+    [
+        ((0, 100), 0, 99),
+        ((0, 25), 0, 24),
+        ((25, 50), 25, 49),
+        ((50, 75), 50, 74),
+        ((75, 100), 75, 99),
+        ((0, 33), 0, 32),
+        ((33, 66), 33, 65),
+        ((66, 100), 66, 99),
+    ],
+)
+async def test_urls_slicing(
+    slice_percent, expected_lower, expected_upper, mock_responses
+):
+    server_url = "http://fake.local/v1"
+    mock_responses.get(
+        server_url + "/",
+        payload={"capabilities": {"attachments": {"base_url": "http://cdn/"}}},
+    )
+    changes_url = server_url + RECORDS_URL.format("monitor", "changes")
+    mock_responses.get(
+        changes_url,
+        payload={
+            "data": [
+                {"id": "abc", "bucket": "bid", "collection": "cid", "last_modified": 42}
+            ]
+        },
+    )
+    records_url = server_url + RECORDS_URL.format("bid", "cid") + "?_expected=42"
+    mock_responses.get(
+        records_url,
+        payload={
+            "data": [
+                {"id": f"id{i}", "attachment": {"location": f"file{i}.jpg"}}
+                for i in range(100)
+            ]
+        },
+    )
+
+    with mock.patch(
+        "checks.remotesettings.attachments_availability.test_url"
+    ) as mocked:
+        await run(server_url, slice_percent=slice_percent)
+    calls = mocked.call_args_list
+    assert calls[0][0] == (f"http://cdn/file{expected_lower}.jpg",)
+    assert calls[-1][0] == (f"http://cdn/file{expected_upper}.jpg",)