Draf: Use different asapo producers for commissioning

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,

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")