feat: stream on-demand inventory queries

tests/test_daemon.py

@@ -1,9 +1,11 @@
 import asyncio
 import os
 import tempfile
+import threading
 import unittest
 from datetime import datetime, timezone
 from pathlib import Path, PurePosixPath
+from unittest.mock import Mock
 from uuid import uuid4
 
 from websockets.asyncio.server import serve
@@ -13,10 +15,67 @@ from archive_clients.daemon import ArchiveClientDaemon
 from archive_clients.probes import FilesystemProbe
 from archive_clients.services import ServiceProbe
 from archive_clients.protocol import decode, encode, encode_message, new_envelope
-from archive_control.v1 import client_pb2, common_pb2, control_pb2, job_pb2
+from archive_control.v1 import (
+    client_pb2, common_pb2, control_pb2, inventory_pb2, job_pb2,
+)
 
 
 class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
+    async def test_slow_inventory_does_not_block_heartbeat(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            token = root / "token"
+            token.write_text("shared-secret", encoding="utf-8")
+            os.chmod(token, 0o600)
+            service = ServiceConfig(
+                "http://local", PurePosixPath("/api"), root
+            )
+            config = ClientConfig(
+                "cache-1", "Cache 1", "cache", "ws://control", token,
+                root / "state.db", root / "backups", service, service,
+            )
+            gate = threading.Event()
+
+            def slow_inventory(_filter=""):
+                gate.wait(2)
+                return []
+
+            reader = Mock(
+                list_resources=Mock(side_effect=slow_inventory),
+                get_resource=Mock(return_value=None),
+            )
+            probe = FilesystemProbe(root, True, True, True, True, True)
+            daemon = ArchiveClientDaemon(
+                config, [probe, probe], [], resource_reader=reader
+            )
+            await asyncio.to_thread(daemon.store.initialize)
+            outbound = asyncio.Queue()
+            tasks = set()
+            command = new_envelope()
+            command.command.command_id = str(uuid4())
+            command.command.created_at.CopyFrom(command.sent_at)
+            command.command.inventory_query.query_id = str(uuid4())
+            command.command.inventory_query.scope = (
+                inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES
+            )
+            await daemon._handle(command, outbound, tasks)
+            self.assertEqual(
+                decode(await outbound.get()).command_ack.status,
+                control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
+            )
+            inventory_task = next(iter(tasks))
+            heartbeat = new_envelope()
+            heartbeat.heartbeat.sequence = 9
+            await daemon._handle(heartbeat, outbound, tasks)
+            self.assertEqual(
+                decode(await outbound.get()).heartbeat_ack.sequence, 9
+            )
+            gate.set()
+            await inventory_task
+            self.assertTrue(
+                decode(await outbound.get()).inventory_chunk.last_chunk
+            )
+
     async def test_registration_heartbeat_and_duplicate_command(self):
         observed = {}
         job_id = str(uuid4())
@@ -39,6 +98,9 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
                 item.service
                 for item in registration.register_request.capabilities.services
             ]
+            observed["features"] = list(
+                registration.register_request.capabilities.features
+            )
             response = new_envelope()
             response.correlation_id = registration.message_id
             response.register_response.status = client_pb2.REGISTRATION_STATUS_ACCEPTED
@@ -80,6 +142,21 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
             observed["rejected_duplicate"] = (
                 decode(await websocket.recv()).command_ack.status
             )
+            inventory = new_envelope()
+            inventory.command.command_id = str(uuid4())
+            inventory.command.created_at.CopyFrom(inventory.sent_at)
+            inventory.command.inventory_query.query_id = str(uuid4())
+            inventory.command.inventory_query.scope = (
+                inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES
+            )
+            await websocket.send(encode(inventory))
+            observed["inventory_ack"] = (
+                decode(await websocket.recv()).command_ack.status
+            )
+            result = decode(await websocket.recv()).inventory_chunk
+            observed["inventory_result"] = (
+                result.last_chunk, result.WhichOneof("payload")
+            )
 
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
@@ -104,7 +181,11 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
                     datetime.now(timezone.utc), version="v2", device_id="DEVICE",
                 )
                 daemon = ArchiveClientDaemon(
-                    config, [probe, probe], [service_probe]
+                    config, [probe, probe], [service_probe],
+                    resource_reader=Mock(
+                        list_resources=Mock(return_value=[]),
+                        get_resource=Mock(return_value=None),
+                    ),
                 )
                 await asyncio.to_thread(daemon.store.initialize)
                 definition = job_pb2.JobDefinition(
@@ -128,6 +209,10 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(observed["addresses"], ["dynamic"])
         self.assertEqual(observed["device_id"], "DEVICE")
         self.assertEqual(observed["services"], ["syncthing"])
+        self.assertIn(
+            client_pb2.CLIENT_FEATURE_INVENTORY_CHUNKS,
+            observed["features"],
+        )
         self.assertEqual(observed["heartbeat"], 7)
         self.assertEqual(observed["first"], control_pb2.COMMAND_ACK_STATUS_ACCEPTED)
         self.assertEqual(observed["second"], control_pb2.COMMAND_ACK_STATUS_DUPLICATE)
@@ -145,6 +230,13 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
             observed["rejected_duplicate"],
             control_pb2.COMMAND_ACK_STATUS_REJECTED,
         )
+        self.assertEqual(
+            observed["inventory_ack"],
+            control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
+        )
+        self.assertEqual(
+            observed["inventory_result"], (True, "resource_summaries")
+        )
 
 
 if __name__ == "__main__":