Daemon.started() is now a context manager

Fixes saltstack#22 Signed-off-by: Pedro Algarvio <palgarvio@vmware.com>
s0undt3ch · May 26, 2022 · 534adbd · 534adbd
1 parent 5d56fb3
commit 534adbd
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 14 deletions.
diff --git a/changelog/22.improvement.rst b/changelog/22.improvement.rst
@@ -0,0 +1 @@
+``Daemon.started()`` is now a context manager
diff --git a/src/pytestshellutils/downgraded/shell.py b/src/pytestshellutils/downgraded/shell.py
@@ -968,21 +968,25 @@ def start(
             start_timeout=start_timeout,
         )
 
+    @contextlib.contextmanager
     def started(
         self,
         *extra_cli_arguments: str,
         max_start_attempts: Optional[int] = None,
         start_timeout: Optional[Union[int, float]] = None
-    ) -> 'Daemon':
+    ) -> Generator['Daemon', None, None]:
         """
         Start the daemon and return it's instance so it can be used as a context manager.
         """
-        self.start(
-            *extra_cli_arguments,
-            max_start_attempts=max_start_attempts,
-            start_timeout=start_timeout,
-        )
-        return self
+        try:
+            self.start(
+                *extra_cli_arguments,
+                max_start_attempts=max_start_attempts,
+                start_timeout=start_timeout,
+            )
+            yield self
+        finally:
+            self.terminate()
 
     @contextlib.contextmanager
     def stopped(
@@ -1061,7 +1065,7 @@ def stopped(
                         exc,
                         exc_info=True,
                     )
-            _started = self.started(*start_arguments.args, **start_arguments.kwargs)
+            _started = self.start(*start_arguments.args, **start_arguments.kwargs)
             if _started:
                 if after_start_callback:
                     try:

diff --git a/src/pytestshellutils/shell.py b/src/pytestshellutils/shell.py
@@ -991,19 +991,25 @@ def start(
             *extra_cli_arguments, max_start_attempts=max_start_attempts, start_timeout=start_timeout
         )
 
+    @contextlib.contextmanager
     def started(
         self,
         *extra_cli_arguments: str,
         max_start_attempts: Optional[int] = None,
         start_timeout: Optional[Union[int, float]] = None
-    ) -> "Daemon":
+    ) -> Generator["Daemon", None, None]:
         """
         Start the daemon and return it's instance so it can be used as a context manager.
         """
-        self.start(
-            *extra_cli_arguments, max_start_attempts=max_start_attempts, start_timeout=start_timeout
-        )
-        return self
+        try:
+            self.start(
+                *extra_cli_arguments,
+                max_start_attempts=max_start_attempts,
+                start_timeout=start_timeout,
+            )
+            yield self
+        finally:
+            self.terminate()
 
     @contextlib.contextmanager
     def stopped(
@@ -1082,7 +1088,7 @@ def stopped(
                         exc,
                         exc_info=True,
                     )
-            _started = self.started(
+            _started = self.start(
                 *start_arguments.args,  # pylint: disable=not-an-iterable
                 **start_arguments.kwargs,  # pylint: disable=not-a-mapping
             )