From e98d25cf850c7c507f42466c215fcfeb346838cd Mon Sep 17 00:00:00 2001
From: Henry Pinkard <7969470+henrypinkard@users.noreply.github.com>
Date: Wed, 6 Sep 2023 22:07:50 -0700
Subject: [PATCH] improve headless shutdown in test

---
 pycromanager/test/conftest.py | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/pycromanager/test/conftest.py b/pycromanager/test/conftest.py
index 41d20cc0..468e041b 100644
--- a/pycromanager/test/conftest.py
+++ b/pycromanager/test/conftest.py
@@ -178,19 +178,18 @@ def setup_data_folder():
 def launch_mm_headless(install_mm):
     mm_install_dir = install_mm
     if mm_install_dir is None:
-        return # local manual testing where MM has been launched from source
-
-    config_file = os.path.join(mm_install_dir, 'MMConfig_demo.cfg')
+        yield # local manual testing where MM has been launched from source
+    else:
+        config_file = os.path.join(mm_install_dir, 'MMConfig_demo.cfg')
+        print('Launching Micro-manager in headless mode.')
 
-    print('Launching Micro-manager in headless mode.')
-    
-    # MM doesn't ship with Java on Mac so allow it to be defined here
-    java_loc = None
-    if "JAVA" in os.environ:
-        java_loc = os.environ["JAVA"]
+        # MM doesn't ship with Java on Mac so allow it to be defined here
+        java_loc = None
+        if "JAVA" in os.environ:
+            java_loc = os.environ["JAVA"]
 
-    start_headless(mm_install_dir, config_file, java_loc=java_loc, debug=True)
+        start_headless(mm_install_dir, config_file, java_loc=java_loc, debug=True)
 
-    yield
+        yield None
 
-    stop_headless()
+        stop_headless()