feat: stream on-demand inventory queries

2026-07-23 01:46:58 +00:00
parent 720fa67202
commit 4058b6d3c8
6 changed files with 568 additions and 9 deletions
@@ -0,0 +1,106 @@
+import unittest
+from uuid import uuid4
+
+from archive_clients.bencode import Metainfo
+from archive_clients.inventory import InventoryService
+from archive_clients.resources import NormalizedResource
+from archive_control.v1 import common_pb2, inventory_pb2, resource_pb2
+
+
+class _Reader:
+    def __init__(self, resources):
+        self.resources = resources
+
+    def list_resources(self, name_filter=""):
+        return [
+            item for item in self.resources
+            if name_filter.casefold() in item.summary.display_name.casefold()
+        ]
+
+    def get_resource(self, torrent_hash):
+        for item in self.resources:
+            identity = item.summary.resource_id
+            if torrent_hash in {
+                identity.info_hash_v1_hex, identity.info_hash_v2_hex,
+            }:
+                return item
+        return None
+
+
+def _resource(index: int) -> NormalizedResource:
+    info_hash = f"{index + 1:040x}"
+    summary = resource_pb2.ResourceSummary(
+        qb_torrent_id=info_hash,
+        display_name=f"Resource {index}",
+        content_revision=f"revision-{index}",
+        canonical_paths=True,
+    )
+    summary.resource_id.info_hash_v1_hex = info_hash
+    summary.selected_files.ranges.add(first=0, last=0)
+    summary.selected_complete_files.ranges.add(first=0, last=0)
+    file = resource_pb2.TorrentFile(
+        file_index=0, canonical_path=f"resource-{index}/file.bin",
+        logical_bytes=10, completed_bytes=10, selected=True,
+    )
+    return NormalizedResource(
+        summary, (file,), Metainfo(info_hash, "", ()),
+    )
+
+
+class InventoryTests(unittest.TestCase):
+    def test_summary_chunks_share_atomic_snapshot_and_page(self):
+        service = InventoryService(
+            _Reader([_resource(index) for index in range(4)]),
+            "cache-1", chunk_target_bytes=1,
+        )
+        query = inventory_pb2.InventoryQuery(
+            query_id=str(uuid4()),
+            scope=inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES,
+        )
+        query.page.page_size = 3
+        chunks = service.execute(query)
+        self.assertEqual(len(chunks), 3)
+        self.assertEqual([chunk.chunk_index for chunk in chunks], [0, 1, 2])
+        self.assertEqual(len({chunk.snapshot_id for chunk in chunks}), 1)
+        self.assertFalse(chunks[0].last_chunk)
+        self.assertTrue(chunks[-1].last_chunk)
+        self.assertEqual(chunks[-1].next_page_token, "3")
+        query.expected_revision = chunks[-1].revision
+        query.page.page_token = chunks[-1].next_page_token
+        next_page = service.execute(query)
+        self.assertFalse(next_page[0].HasField("error"))
+        self.assertEqual(
+            next_page[0].resource_summaries.resources[0].display_name,
+            "Resource 3",
+        )
+
+    def test_content_tree_requires_matching_revision(self):
+        resource = _resource(0)
+        service = InventoryService(_Reader([resource]), "cache-1")
+        query = inventory_pb2.InventoryQuery(
+            query_id=str(uuid4()),
+            scope=inventory_pb2.INVENTORY_SCOPE_CONTENT_TREE,
+            expected_revision="stale",
+        )
+        query.resource_ids.add().CopyFrom(resource.summary.resource_id)
+        stale = service.execute(query)
+        self.assertEqual(stale[0].error.code, common_pb2.ERROR_CODE_STALE_STATE)
+        query.expected_revision = resource.summary.content_revision
+        chunks = service.execute(query)
+        self.assertEqual(chunks[0].WhichOneof("payload"), "content_tree")
+        self.assertEqual(chunks[0].content_tree.entries[0].available_file_count, 1)
+
+    def test_invalid_query_is_a_terminal_error_chunk(self):
+        query = inventory_pb2.InventoryQuery(
+            query_id="invalid",
+            scope=inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES,
+        )
+        result = InventoryService(_Reader([]), "cache-1").execute(query)
+        self.assertTrue(result[0].last_chunk)
+        self.assertEqual(
+            result[0].error.code, common_pb2.ERROR_CODE_INVALID_ARGUMENT
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()