From fc00824969967d6bc00f43755855436647e33073 Mon Sep 17 00:00:00 2001 From: Pouria Khalaj Date: Wed, 29 May 2024 11:02:59 +0200 Subject: [PATCH] Refactor `cli` and `test_recorder` --- trolldb/cli.py | 24 ++++++++++++++---------- trolldb/tests/test_recorder.py | 9 ++++----- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/trolldb/cli.py b/trolldb/cli.py index d5e1a01..55830a6 100644 --- a/trolldb/cli.py +++ b/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): diff --git a/trolldb/tests/test_recorder.py b/trolldb/tests/test_recorder.py index d54bb6e..512c719 100644 --- a/trolldb/tests/test_recorder.py +++ b/trolldb/tests/test_recorder.py @@ -101,13 +101,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 message_in_database_and_delete_count_is_one(msg)