107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
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()
|