tests: always shutdown server, even in failures (#3837)

marimo-team · Feb 18, 2025 · 58ed6c0 · 58ed6c0
1 parent 8a3ab96
commit 58ed6c0
Showing 1 changed file with 51 additions and 48 deletions.
diff --git a/tests/_server/mocks.py b/tests/_server/mocks.py
@@ -99,25 +99,26 @@ def wrapper(
             auth_token = get_session_manager(client).auth_token
             headers = token_header(auth_token)
 
-            with client.websocket_connect(
-                f"/ws?session_id={session_id}", headers=headers
-            ) as websocket:
-                data = websocket.receive_text()
-                assert data
-                if "temp_marimo_file" in func.__code__.co_varnames:
-                    func(
-                        client,
-                        temp_marimo_file=temp_marimo_file,
+            try:
+                with client.websocket_connect(
+                    f"/ws?session_id={session_id}", headers=headers
+                ) as websocket:
+                    data = websocket.receive_text()
+                    assert data
+                    if "temp_marimo_file" in func.__code__.co_varnames:
+                        func(
+                            client,
+                            temp_marimo_file=temp_marimo_file,
+                        )
+                    else:
+                        func(client)
+            finally:
+                # Always shutdown, even if there's an error
+                if auto_shutdown:
+                    client.post(
+                        "/api/kernel/shutdown",
+                        headers=headers,
                     )
-                else:
-                    func(client)
-            # shutdown after websocket exits, otherwise
-            # test fails on Windows (loop closed twice)
-            if auto_shutdown:
-                client.post(
-                    "/api/kernel/shutdown",
-                    headers=headers,
-                )
 
         return wrapper
 
@@ -135,20 +136,21 @@ def wrapper(client: TestClient) -> None:
             auth_token = get_session_manager(client).auth_token
             headers = token_header(auth_token)
 
-            with client.websocket_connect(
-                f"/ws?session_id={session_id}", headers=headers
-            ) as websocket:
-                data = websocket.receive_text()
-                assert data
-
-                func(client, websocket)
-            # shutdown after websocket exits, otherwise
-            # test fails on Windows (loop closed twice)
-            if auto_shutdown:
-                client.post(
-                    "/api/kernel/shutdown",
-                    headers=headers,
-                )
+            try:
+                with client.websocket_connect(
+                    f"/ws?session_id={session_id}", headers=headers
+                ) as websocket:
+                    data = websocket.receive_text()
+                    assert data
+
+                    func(client, websocket)
+            finally:
+                # Always shutdown, even if there's an error
+                if auto_shutdown:
+                    client.post(
+                        "/api/kernel/shutdown",
+                        headers=headers,
+                    )
 
         return wrapper
 
@@ -165,22 +167,23 @@ def wrapper(client: TestClient) -> None:
             session_manager = get_session_manager(client)
             headers = token_header(session_manager.auth_token)
 
-            with client.websocket_connect(
-                f"/ws?session_id={session_id}", headers=headers
-            ) as websocket:
-                data = websocket.receive_text()
-                assert data
-                # Just change the mode here, otherwise our tests will run,
-                # in threads
-                session_manager.mode = SessionMode.RUN
-                func(client)
-                session_manager.mode = SessionMode.EDIT
-            # shutdown after websocket exits, otherwise
-            # test fails on Windows (loop closed twice)
-            client.post(
-                "/api/kernel/shutdown",
-                headers=headers,
-            )
+            try:
+                with client.websocket_connect(
+                    f"/ws?session_id={session_id}", headers=headers
+                ) as websocket:
+                    data = websocket.receive_text()
+                    assert data
+                    # Just change the mode here, otherwise our tests will run,
+                    # in threads
+                    session_manager.mode = SessionMode.RUN
+                    func(client)
+                    session_manager.mode = SessionMode.EDIT
+            finally:
+                # Always shutdown, even if there's an error
+                client.post(
+                    "/api/kernel/shutdown",
+                    headers=headers,
+                )
 
         return wrapper