Refactor cli and test_recorder

trolldb/cli.py

@@ -20,16 +20,20 @@ async def record_messages(config: AppConfig):
         )
         for m in create_subscriber_from_dict_config(config.subscriber).recv():
             msg = Message.decode(str(m))
-            if msg.type in ["file", "dataset"]:
-                await collection.insert_one(msg.data)
-                logger.info(f"Inserted file with uri: {msg.data["uri"]}")
-            elif msg.type == "del":
-                deletion_result = await collection.delete_many({"uri": msg.data["uri"]})
-                logger.info(f"Deleted document with uri: {msg.data["uri"]}")
-                if deletion_result.deleted_count != 1:
-                    logger.error(f"Recorder found multiple deletions for uri: {msg.data["uri"]}!")
-            else:
-                logger.debug(f"Don't know what to do with {msg.type} message.")
+            match msg.type:
+                case "file":
+                    await collection.insert_one(msg.data)
+                    logger.info(f"Inserted file with uri: {msg.data["uri"]}")
+                case "dataset":
+                    await collection.insert_one(msg.data)
+                    logger.info(f"Inserted dataset with {len(msg.data["dataset"])} elements.")
+                case "del":
+                    deletion_result = await collection.delete_many({"uri": msg.data["uri"]})
+                    logger.info(f"Deleted document with uri: {msg.data["uri"]}")
+                    if deletion_result.deleted_count != 1:
+                        logger.error(f"Recorder found multiple deletions for uri: {msg.data["uri"]}!")
+                case _:
+                    logger.debug(f"Don't know what to do with {msg.type} message.")
 
 
 async def record_messages_from_config(config_file: FilePath):

trolldb/tests/test_recorder.py

@@ -57,13 +57,16 @@ def config_file(tmp_path):
     return create_config_file(tmp_path)
 
 
-async def message_in_database_and_delete_count_is_one(msg) -> bool:
+async def message_in_database_and_delete_count_is_one(msg, is_dataset_message: bool = False) -> bool:
     """Checks if there is exactly one item in the database which matches the data of the message."""
     async with mongodb_context(test_app_config.database):
         collection = await MongoDB.get_collection("test_database", "test_collection")
         result = await collection.find_one(dict(scan_mode="EW"))
         result.pop("_id")
-        deletion_result = await collection.delete_many({"uri": msg.data["uri"]})
+        if is_dataset_message:
+            deletion_result = await collection.delete_many({"dataset.uri": msg.data["dataset"][0]["uri"]})
+        else:
+            deletion_result = await collection.delete_many({"uri": msg.data["uri"]})
         return result == msg.data and deletion_result.deleted_count == 1
 
 
@@ -101,13 +104,12 @@ async def test_record_deletes_message(tmp_path, file_message, del_message):
                 result = await collection.find_one(dict(scan_mode="EW"))
                 assert result is None
 
+
 async def test_record_dataset_messages(tmp_path, dataset_message):
-    """Test recording a dataset message and deleting the file."""
+    """Tests recording a dataset message and deleting the file."""
     config = AppConfig(**make_test_app_config(tmp_path))
+    msg = Message.decode(dataset_message)
     with running_prepared_database_context():
         with patched_subscriber_recv([dataset_message]):
             await record_messages(config)
-            async with mongodb_context(config.database):
-                collection = await MongoDB.get_collection("mock_database", "mock_collection")
-                result = await collection.find_one(dict(scan_mode="EW"))
-                assert result is not None
+            assert await message_in_database_and_delete_count_is_one(msg, is_dataset_message=True)