diff --git a/src/hidra/receiver/plugins/asapo_producer.py b/src/hidra/receiver/plugins/asapo_producer.py index 5c64f06e..c373f4c6 100644 --- a/src/hidra/receiver/plugins/asapo_producer.py +++ b/src/hidra/receiver/plugins/asapo_producer.py @@ -75,6 +75,21 @@ logger = logging.getLogger(__name__) +def get_beamtime(metadata): + path_parts = Path(metadata["relative_path"]).parts + if path_parts[0] == "current": + beamtime = "auto" + elif path_parts[0] == "commissioning": + beamtime = "commissioning" + else: + raise utils.NotSupported( + "Path '{}' is not supported".format( + Path().joinpath(*path_parts).as_posix()) + ) + + return beamtime + + def get_exposed_path(metadata): exposed_path = Path(metadata["relative_path"], metadata["filename"]).parts @@ -271,11 +286,13 @@ def __init__(self, endpoint, beamtime, token, n_threads, file_regex, # degraded performance as producers would be constantly deleted and recreated. self.max_active_data_sources = 10 - def _create_producer(self, data_source): - logger.info("Create producer with data_source=%s", data_source) - self.data_source_info[data_source] = { + def _create_producer(self, beamtime, data_source): + logger.info( + "Create producer with beamtime=%s data_source=%s", beamtime, data_source + ) + self.data_source_info[(beamtime, data_source)] = { "producer": asapo_producer.create_producer( - self.endpoint, "raw", self.beamtime, self.beamline, + self.endpoint, "raw", beamtime, self.beamline, data_source, self.token, self.n_threads, self.timeout * 1000), } @@ -299,13 +316,13 @@ def _create_producer(self, data_source): " for data_source=%s", oldest_data_source) oldest_producer.cleanup() - def _get_producer(self, data_source): - if data_source not in self.data_source_info: - self._create_producer(data_source=data_source) + def _get_producer(self, beamtime, data_source): + if (beamtime, data_source) not in self.data_source_info: + self._create_producer(beamtime=beamtime, data_source=data_source) else: # Move most recently used producers to the end - self.data_source_info.move_to_end(data_source) - return self.data_source_info[data_source]["producer"] + self.data_source_info.move_to_end((beamtime, data_source)) + return self.data_source_info[(beamtime, data_source)]["producer"] def send_message(self, local_path, metadata): try: @@ -320,7 +337,12 @@ def send_message(self, local_path, metadata): logger.debug("Ignoring file %s", local_path) return - producer = self._get_producer(data_source) + if self.beamtime == "auto": + beamtime = get_beamtime(metadata) + else: + beamtime = self.beamtime + + producer = self._get_producer(beamtime, data_source) producer.send( # files start with index 0 and asapo with 1 id=file_idx + 1 - self.start_file_idx, diff --git a/test/pytest/receiver/plugins/test_asapo_producer.py b/test/pytest/receiver/plugins/test_asapo_producer.py index 611dbcbd..cfe283d4 100644 --- a/test/pytest/receiver/plugins/test_asapo_producer.py +++ b/test/pytest/receiver/plugins/test_asapo_producer.py @@ -13,7 +13,8 @@ def config(): config = dict( endpoint="asapo-services:8400", - beamtime="p00", + beamtime="auto", + beamline="p00", token="abcdefg1234=", default_data_source='test001', n_threads=1, @@ -80,17 +81,45 @@ def test_worker_create_producer_commissioning(worker, mock_create_producer, conf } worker.send_message(filepath, metadata) + mock_create_producer.assert_called_once_with( + config["endpoint"], + 'raw', + "commissioning", + "p00", + config["default_data_source"], + config["token"], + config["n_threads"], + config.get("timeout", 5) * 1000, + ) + + +def test_worker_create_producer_fixed_beamtime(config, mock_create_producer): + config["beamtime"] = "abc012" + del config["beamline"] + del config["user_config_path"] + + worker = AsapoWorker(**config) + + filepath = "/tmp/hidra_source/current/raw/det01/stream100_scan0-107.tif" + metadata = { + "relative_path": "current/raw/det01", + "filename": "stream100_scan0-107.tif" + } + worker.send_message(filepath, metadata) + mock_create_producer.assert_called_once_with( config["endpoint"], 'raw', config["beamtime"], - config.get("beamline", "auto"), + "auto", config["default_data_source"], config["token"], config["n_threads"], config.get("timeout", 5) * 1000, ) + assert (config["beamtime"], config["default_data_source"]) in worker.data_source_info + def test_worker_send_message(worker, mock_producer): filepath = "/tmp/hidra_source/current/raw/det01/stream100_scan0-107.tif" @@ -235,6 +264,8 @@ def test_worker_data_source_removal(config, mock_create_producer): ) worker = AsapoWorker(**config) + beamtime = config["beamtime"] + # Create 11 data sources/producers for i in range(worker.max_active_data_sources + 1): filepath = "/tmp/hidra_source/current/raw/det{:02d}/stream100_scan0-107.tif".format(i) @@ -246,8 +277,8 @@ def test_worker_data_source_removal(config, mock_create_producer): assert "det{:02d}".format(i) in mock_create_producer.call_args.args assert len(worker.data_source_info) == worker.max_active_data_sources - assert "det00" not in worker.data_source_info - assert next(iter(worker.data_source_info)) == "det01" + assert (beamtime, "det00") not in worker.data_source_info + assert next(iter(worker.data_source_info)) == (beamtime, "det01") # Send another message to move data source det01 to the end filepath = ( @@ -259,7 +290,7 @@ def test_worker_data_source_removal(config, mock_create_producer): } worker.send_message(filepath, metadata) - assert next(iter(worker.data_source_info)) == "det02" + assert next(iter(worker.data_source_info)) == (beamtime, "det02") # Create a new data source filepath = ( @@ -273,8 +304,8 @@ def test_worker_data_source_removal(config, mock_create_producer): assert "det99" in mock_create_producer.call_args.args assert len(worker.data_source_info) == worker.max_active_data_sources - assert "det00" not in worker.data_source_info - assert "det02" not in worker.data_source_info - assert "det01" in worker.data_source_info - assert "det99" in worker.data_source_info - assert next(iter(worker.data_source_info)) == "det03" + assert (beamtime, "det00") not in worker.data_source_info + assert (beamtime, "det02") not in worker.data_source_info + assert (beamtime, "det01") in worker.data_source_info + assert (beamtime, "det99") in worker.data_source_info + assert next(iter(worker.data_source_info)) == (beamtime, "det03")