From ba068a0f5822dccddb134b8fc520d1716c4ed1f7 Mon Sep 17 00:00:00 2001
From: Justin Newberry <justin@comma.ai>
Date: Wed, 6 Mar 2024 18:22:07 -0500
Subject: [PATCH] with_processes: add standalone context manager (#31746)

helpers
---
 selfdrive/test/helpers.py | 40 +++++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/selfdrive/test/helpers.py b/selfdrive/test/helpers.py
index 210a283699..b345b929ec 100644
--- a/selfdrive/test/helpers.py
+++ b/selfdrive/test/helpers.py
@@ -1,3 +1,4 @@
+import contextlib
 import http.server
 import os
 import threading
@@ -43,27 +44,34 @@ def release_only(f):
     f(self, *args, **kwargs)
   return wrap
 
-def with_processes(processes, init_time=0, ignore_stopped=None):
+
+@contextlib.contextmanager
+def processes_context(processes, init_time=0, ignore_stopped=None):
   ignore_stopped = [] if ignore_stopped is None else ignore_stopped
 
+  # start and assert started
+  for n, p in enumerate(processes):
+    managed_processes[p].start()
+    if n < len(processes) - 1:
+      time.sleep(init_time)
+
+  assert all(managed_processes[name].proc.exitcode is None for name in processes)
+
+  try:
+    yield [managed_processes[name] for name in processes]
+    # assert processes are still started
+    assert all(managed_processes[name].proc.exitcode is None for name in processes if name not in ignore_stopped)
+  finally:
+    for p in processes:
+      managed_processes[p].stop()
+
+
+def with_processes(processes, init_time=0, ignore_stopped=None):
   def wrapper(func):
     @wraps(func)
     def wrap(*args, **kwargs):
-      # start and assert started
-      for n, p in enumerate(processes):
-        managed_processes[p].start()
-        if n < len(processes) - 1:
-          time.sleep(init_time)
-      assert all(managed_processes[name].proc.exitcode is None for name in processes)
-
-      # call the function
-      try:
-        func(*args, **kwargs)
-        # assert processes are still started
-        assert all(managed_processes[name].proc.exitcode is None for name in processes if name not in ignore_stopped)
-      finally:
-        for p in processes:
-          managed_processes[p].stop()
+      with processes_context(processes, init_time, ignore_stopped):
+        return func(*args, **kwargs)
 
     return wrap
   return wrapper