fix(bigquery-magics): wait for graph server startup before returning from init

packages/bigquery-magics/bigquery_magics/graph_server.py

@@ -297,6 +297,7 @@ def __init__(self):
         self.port = None
         self.url = None
         self._server = None
+        self._server_ready = threading.Event()
 
     def build_route(self, endpoint):
         """
@@ -317,6 +318,7 @@ class ThreadedTCPServer(socketserver.TCPServer):
 
         with ThreadedTCPServer(("127.0.0.1", self.port), GraphServerHandler) as httpd:
             self._server = httpd
+            self._server_ready.set()
             self._server.serve_forever()
 
     def init(self):
@@ -327,9 +329,11 @@ def init(self):
 
         self.port = portpicker.pick_unused_port()
         self.url = f"{GraphServer.host}:{self.port}"
+        self._server_ready.clear()
 
         server_thread = threading.Thread(target=self._start_server)
         server_thread.start()
+        self._server_ready.wait()
         return server_thread
 
     def stop_server(self):

packages/bigquery-magics/tests/unit/test_graph_server.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import time
 import unittest
 from unittest import mock
 
@@ -747,3 +748,55 @@ def test_convert_schema_shared_label():
     labels = {label["name"]: label for label in result["labels"]}
     assert "Person" in labels
     assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"}
+
+
+def test_stop_server_immediately_after_init():
+    """Verifies that GraphServer.init() waits until the underlying TCP server
+    is initialized before returning, preventing a race condition with stop_server().
+    """
+    # Mock portpicker to return port 0 (OS-assigned ephemeral port) so this test
+    # runs in all unit test environments, even when optional dependencies like
+    # `spanner-graph-notebook` and `portpicker` are not installed.
+    mock_portpicker = mock.MagicMock()
+    mock_portpicker.pick_unused_port.return_value = 0
+
+    with mock.patch.dict("sys.modules", {"portpicker": mock_portpicker}):
+        server = graph_server.GraphServer()
+        server_thread = server.init()
+
+        try:
+            # 1. Immediately call stop_server() after init() without making any HTTP
+            #    requests. If init() returns before `self._server` is assigned in
+            #    _start_server(), `if self._server:` inside stop_server() evaluates to
+            #    False, making stop_server() a no-op.
+            server.stop_server()
+
+            # 2. Wait up to 1.0 second for the background server thread to exit.
+            #    When properly synchronized, init() blocks until `self._server` is
+            #    assigned, so stop_server() calls `self._server.shutdown()` and the
+            #    thread exits almost immediately. If the shutdown signal is missed,
+            #    the background thread remains blocked in `serve_forever()`. Using
+            #    `timeout=1.0` prevents `join()` from hanging indefinitely.
+            server_thread.join(timeout=1.0)
+
+            # 3. Assert that the server thread actually terminated.
+            assert not server_thread.is_alive(), (
+                "GraphServer thread is still alive after stop_server(); "
+                "stop_server() likely raced with _start_server() before self._server was set."
+            )
+        finally:
+            # 4. Cleanup in case the assertion above fails:
+            #    Because `server_thread` is a non-daemon thread, Python's interpreter
+            #    shutdown (`threading._shutdown()`) will hang at the end of the pytest
+            #    run if `server_thread` is left running.
+            if server_thread.is_alive():  # pragma: NO COVER
+                # Poll with a short `time.sleep()` (releasing the GIL so the background
+                # thread can finish binding the socket and setting `server._server`)
+                # up to a bounded deadline.
+                deadline = time.monotonic() + 5.0
+                while server._server is None and time.monotonic() < deadline:
+                    time.sleep(0.01)
+
+                # Shut down the server so the background thread exits cleanly.
+                server.stop_server()
+                server_thread.join(timeout=2.0)