From fe8e2c05fd32b706230d02f40cd752e5ee1faf46 Mon Sep 17 00:00:00 2001
From: Justin Newberry <justin@comma.ai>
Date: Thu, 7 Mar 2024 15:09:07 -0500
Subject: [PATCH] test helpers: http server context (#31774)

* http context

* fix

* fix
old-commit-hash: 8d9b96cf230830b3ce88672e383b4b0e50f7f887
---
 selfdrive/test/helpers.py | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/selfdrive/test/helpers.py b/selfdrive/test/helpers.py
index b345b929ec..c157b98b62 100644
--- a/selfdrive/test/helpers.py
+++ b/selfdrive/test/helpers.py
@@ -88,23 +88,28 @@ def read_segment_list(segment_list_path):
   return [(platform[2:], segment) for platform, segment in zip(seg_list[::2], seg_list[1::2], strict=True)]
 
 
+@contextlib.contextmanager
+def http_server_context(handler, setup=None):
+  host = '127.0.0.1'
+  server = http.server.HTTPServer((host, 0), handler)
+  port = server.server_port
+  t = threading.Thread(target=server.serve_forever)
+  t.start()
+
+  if setup is not None:
+    setup(host, port)
+
+  try:
+    yield (host, port)
+  finally:
+    server.shutdown()
+    server.server_close()
+    t.join()
+
+
 def with_http_server(func, handler=http.server.BaseHTTPRequestHandler, setup=None):
   @wraps(func)
   def inner(*args, **kwargs):
-    host = '127.0.0.1'
-    server = http.server.HTTPServer((host, 0), handler)
-    port = server.server_port
-    t = threading.Thread(target=server.serve_forever)
-    t.start()
-
-    if setup is not None:
-      setup(host, port)
-
-    try:
-      return func(*args, f'http://{host}:{port}', **kwargs)
-    finally:
-      server.shutdown()
-      server.server_close()
-      t.join()
-
+    with http_server_context(handler, setup) as (host, port):
+      return func(*args, f"http://{host}:{port}", **kwargs)
   return inner