From cecf6e2078d8dce1fa2eee059de7e785056ee833 Mon Sep 17 00:00:00 2001
From: ilyakolb <ilya.kolb.0@gmail.com>
Date: Fri, 25 Oct 2024 11:42:13 -0700
Subject: [PATCH] fix metadata return logic, add test

---
 pycromanager/acquisition/acq_future.py          | 14 +++++++-------
 pycromanager/test/test_notifications_futures.py |  8 ++++++++
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/pycromanager/acquisition/acq_future.py b/pycromanager/acquisition/acq_future.py
index 58fd9545..aab06fd5 100644
--- a/pycromanager/acquisition/acq_future.py
+++ b/pycromanager/acquisition/acq_future.py
@@ -139,20 +139,20 @@ def await_image_saved(self, axes=None, return_image=True, return_metadata=False)
                         while not self._notification_recieved[key][AcqNotification.Image.IMAGE_SAVED]:
                             self._condition.wait()
 
-        if return_image:
+        if return_image and return_metadata:
+            if isinstance(axes, list):
+                return [(self._acq.get_dataset().read_image(**ax), self._acq.get_dataset().read_metadata(**ax)) for ax in axes]
+            else:
+                return self._acq.get_dataset().read_image(**axes), self._acq.get_dataset().read_metadata(**axes)
+        elif return_image:
             if isinstance(axes, list):
                 return [self._acq.get_dataset().read_image(**ax) for ax in axes]
             else:
                 return self._acq.get_dataset().read_image(**axes)
-        if return_metadata:
+        elif return_metadata:
             if isinstance(axes, list):
                 return [self._acq.get_dataset().read_metadata(**ax) for ax in axes]
             else:
                 return self._acq.get_dataset().read_metadata(**axes)
-        if return_image and return_metadata:
-            if isinstance(axes, list):
-                return [(self._acq.get_dataset().read_image(**ax), self._acq.get_dataset().read_metadata(**ax)) for ax in axes]
-            else:
-                return self._acq.get_dataset().read_image(**axes), self._acq.get_dataset().read_metadata(**axes)
 
 
diff --git a/pycromanager/test/test_notifications_futures.py b/pycromanager/test/test_notifications_futures.py
index c4e7eae8..e19ce2d3 100644
--- a/pycromanager/test/test_notifications_futures.py
+++ b/pycromanager/test/test_notifications_futures.py
@@ -15,6 +15,14 @@ def test_async_image_read(launch_mm_headless, setup_data_folder):
         assert np.all(image == acq.get_dataset().read_image(time=5))
     acq.get_dataset().close()
 
+def test_async_image_read_metadata_return(launch_mm_headless, setup_data_folder):
+    events = multi_d_acquisition_events(num_time_points=10, time_interval_s=0.5)
+    with Acquisition(directory=setup_data_folder, name='test_async_image_read', show_display=False) as acq:
+        future = acq.acquire(events)
+        _, metadata = future.await_image_saved({'time': 5}, return_image=True, return_metadata=True)
+        assert isinstance(metadata, dict)
+    acq.get_dataset().close()
+
 def test_async_image_read_sequence(launch_mm_headless, setup_data_folder):
     events = multi_d_acquisition_events(num_time_points=10, time_interval_s=0)
     with Acquisition(directory=setup_data_folder, name='test_async_image_read_sequence', show_display=False) as acq: