Move get_ping() and post_ping() out of the GraphServer class, into the unit test

Author: ericfe-google

bigquery_magics/graph_server.py

@@ -168,27 +168,7 @@ def stop_server():
         if GraphServer._server:
             GraphServer._server.shutdown()
             print("BigQuery-magics graph server shutting down...")
-
-    @staticmethod
-    def get_ping():
-        route = GraphServer.build_route(GraphServer.endpoints["get_ping"])
-        response = requests.get(route)
-
-        assert (
-            response.status_code == 200
-        )  # Guaranteed by GraphServerHandler implementation
-        return response.json()
-
-    @staticmethod
-    def post_ping(data):
-        route = GraphServer.build_route(GraphServer.endpoints["post_ping"])
-        response = requests.post(route, json=data)
-
-        if response.status_code == 200:
-            return response.json()
-        else:
-            print(f"Request failed with status code {response.status_code}")
-            return False
+            GraphServer._server = None
 
 
 class GraphServerHandler(http.server.SimpleHTTPRequestHandler):

tests/unit/test_graph_server.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import requests
 import unittest
 
 import pytest
@@ -320,15 +321,20 @@ def tearDown(self):
         GraphServer.stop_server()  # Stop the server after each test
         self.server_thread.join()  # Wait for the thread to finish
 
-    def test_ping(self):
+    def test_get_ping(self):
         self.assertTrue(self.server_thread.is_alive())
 
-        response = GraphServer.get_ping()
-        self.assertEqual(response, {"message": "pong"})
+        route = GraphServer.build_route(GraphServer.endpoints["get_ping"])
+        response = requests.get(route)
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.json(), {"message": "pong"})
 
-        request = {"data": "ping"}
-        response = GraphServer.post_ping(request)
-        self.assertEqual(response, {"your_request": request})
+    def test_post_ping(self):
+        self.assertTrue(self.server_thread.is_alive())
+        route = GraphServer.build_route(GraphServer.endpoints["post_ping"])
+        response = requests.post(route, json={"data": "ping"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.json(), {"your_request": {"data": "ping"}})
 
 
 def test_stop_server_never_started():