From b24bc53305293d227a72a48f0444326a79680634 Mon Sep 17 00:00:00 2001 From: Martin Raspaud Date: Wed, 29 May 2024 10:12:06 +0200 Subject: [PATCH] Add support for "dataset" messages --- trolldb/cli.py | 17 ++++++++--------- trolldb/tests/test_recorder.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/trolldb/cli.py b/trolldb/cli.py index 74ae37b..294aab9 100644 --- a/trolldb/cli.py +++ b/trolldb/cli.py @@ -20,15 +20,14 @@ async def record_messages(config: AppConfig): ) for m in create_subscriber_from_dict_config(config.subscriber).recv(): msg = Message.decode(str(m)) - match msg.type: - case "file": - await collection.insert_one(msg.data) - case "del": - deletion_result = await collection.delete_many({"uri": msg.data["uri"]}) - if deletion_result.deleted_count != 1: - logger.error("Recorder found multiple deletions!") # TODO: Log some data related to the msg - case _: - logger.debug(f"Don't know what to do with {msg.type} message.") + if msg.type in ["file", "dataset"]: + await collection.insert_one(msg.data) + elif msg.type == "del": + deletion_result = await collection.delete_many({"uri": msg.data["uri"]}) + if deletion_result.deleted_count != 1: + logger.error("Recorder found multiple deletions!") # TODO: Log some data related to the msg + else: + 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 cf90da2..99546c6 100644 --- a/trolldb/tests/test_recorder.py +++ b/trolldb/tests/test_recorder.py @@ -22,6 +22,17 @@ def file_message(tmp_data_filename): '"polarization": "hh", "sensor": "sar-c", "format": "GeoTIFF", "pass_direction": "ASCENDING"}') +@pytest.fixture() +def dataset_message(tmp_data_filename): + """Create a string for a file message.""" + return ('pytroll://segment/raster/L2/SAR dataset a001673@c20969.ad.smhi.se 2019-11-05T13:00:10.366023 v1.01 ' + 'application/json {"platform_name": "S1B", "scan_mode": "EW", "type": "GRDM", "data_source": "1SDH", ' + '"start_time": "2019-11-03T15:39:36.543000", "end_time": "2019-11-03T15:40:40.821000", "orbit_number": ' + '18765, "random_string1": "0235EA", "random_string2": "747D", "dataset": [{"uri": ' + f'"{str(tmp_data_filename)}", "uid": "20191103_153936-s1b-ew-hh.tiff"}}], ' + '"polarization": "hh", "sensor": "sar-c", "format": "GeoTIFF", "pass_direction": "ASCENDING"}') + + @pytest.fixture() def del_message(tmp_data_filename): """Create a string for a delete message.""" @@ -89,3 +100,14 @@ async def test_record_deletes_message(tmp_path, file_message, del_message): collection = await MongoDB.get_collection("mock_database", "mock_collection") 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.""" + config = AppConfig(**make_test_app_config(tmp_path)) + 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