Simplify tests for initial version

DiamondLightSource · Jul 31, 2023 · 60354a8 · 60354a8
1 parent 23c72d5
commit 60354a8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/src/blueapi/plugins/data_writing.py b/src/blueapi/plugins/data_writing.py
@@ -69,26 +69,24 @@ def data_writing_wrapper(
         provider = InMemoryDataCollectionProvider()
 
     scan_number = itertools.count()
-    # next_scan_number = None
+    next_scan_number = None
     stage_stack: Deque = deque()
-    scan_number_stack: Deque = deque()
+    # scan_number_stack: Deque = deque()
     for message in plan:
         if message.command == "stage":
             stage_stack.append(message.obj)
         elif stage_stack:
-            scan_number_stack.append(next(scan_number))
-            next_scan_number = scan_number_stack[-1]
+            next_scan_number = next(scan_number)
             root_devices = []
             while stage_stack:
                 root_devices.append(stage_stack.pop())
             all_devices = walk_devices(root_devices)
             collection = provider.get_next_data_collection(collection_group)
             configure_data_writing(all_devices, collection)
 
-        if message.command == "open_run" and "scan_number" not in message.kwargs:
-            if not scan_number_stack:
-                scan_number_stack.append(next(scan_number))
-            next_scan_number = scan_number_stack[-1]
+        if message.command == "open_run":
+            if next_scan_number is None:
+                next_scan_number = next(scan_number)
             message.kwargs["scan_number"] = next_scan_number
         yield message
 

diff --git a/tests/plugins/test_data_writing.py b/tests/plugins/test_data_writing.py
@@ -121,10 +121,10 @@ def test_multi_run_single_stage_multi_group(
     docs = collect_docs(run_engine, multi_run_single_stage_multi_group(detectors))
     start_docs = find_start_docs(docs)
     assert len(start_docs) == 4
-    assert start_docs[0].doc["scan_number"] == 1
-    assert start_docs[1].doc["scan_number"] == 1
-    assert start_docs[2].doc["scan_number"] == 2
-    assert start_docs[3].doc["scan_number"] == 2
+    assert start_docs[0].doc["scan_number"] == 0
+    assert start_docs[1].doc["scan_number"] == 0
+    assert start_docs[2].doc["scan_number"] == 0
+    assert start_docs[3].doc["scan_number"] == 0
 
 
 def test_nested_run_with_metadata(
@@ -134,9 +134,9 @@ def test_nested_run_with_metadata(
     docs = collect_docs(run_engine, nested_run_with_metadata(detectors))
     start_docs = find_start_docs(docs)
     assert len(start_docs) == 3
-    assert start_docs[0].doc["scan_number"] == 12345
-    assert start_docs[1].doc["scan_number"] == 12345
-    assert start_docs[2].doc["scan_number"] == 12345
+    assert start_docs[0].doc["scan_number"] == 0
+    assert start_docs[1].doc["scan_number"] == 1
+    assert start_docs[2].doc["scan_number"] == 2
 
 
 def test_nested_run_without_metadata(
@@ -147,8 +147,8 @@ def test_nested_run_without_metadata(
     start_docs = find_start_docs(docs)
     assert len(start_docs) == 3
     assert start_docs[0].doc["scan_number"] == 0
-    assert start_docs[1].doc["scan_number"] == 0
-    assert start_docs[2].doc["scan_number"] == 0
+    assert start_docs[1].doc["scan_number"] == 1
+    assert start_docs[2].doc["scan_number"] == 2
 
 
 def collect_docs(run_engine: RunEngine, plan: MsgGenerator) -> List[DataEvent]: