Add more tests.

Author: ericfe-google

bigquery_magics/graph_server.py

@@ -174,11 +174,10 @@ def get_ping():
         route = GraphServer.build_route(GraphServer.endpoints["get_ping"])
         response = requests.get(route)
 
-        if response.status_code == 200:
-            return response.json()
-        else:
-            print(f"Request failed with status code {response.status_code}")
-            return False
+        assert (
+            response.status_code == 200
+        )  # Guaranteed by GraphServerHandler implementation
+        return response.json()
 
     @staticmethod
     def post_ping(data):

tests/unit/test_graph_server.py

@@ -296,6 +296,22 @@ def test_convert_empty_dict():
     }
 
 
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_wrong_row_index():
+    result = convert_graph_data(
+        {
+            "result": {
+                # Missing "0" key
+                "1": json.dumps(row_alex_owns_account),
+            }
+        }
+    )
+
+    assert result == {"error": "Unexpected row index; expected 0, got 1"}
+
+
 class TestGraphServer(unittest.TestCase):
     def setUp(self):
         self.server_thread = GraphServer.init()
@@ -313,3 +329,7 @@ def test_ping(self):
         request = {"data": "ping"}
         response = GraphServer.post_ping(request)
         self.assertEqual(response, {"your_request": request})
+
+
+def test_stop_server_never_started():
+    GraphServer.stop_server()