Fix unit tests

Author: ericfe-google

bigquery_magics/graph_server.py

@@ -251,6 +251,7 @@ def handle_post_ping(self):
         self.do_data_response({"your_request": data})
 
     def handle_post_query(self):
+        self.parse_post_data()
         query_results = json.loads(graph_server.query_result.to_json())
         response = convert_graph_data(query_results=query_results)
         self.do_data_response(response)

tests/unit/bigquery/test_bigquery.py

@@ -680,18 +680,10 @@ def test_bigquery_graph_json_result(monkeypatch):
         html_content = display_mock.call_args_list[0][0][0].data
         assert "<script>" in html_content
         assert "</script>" in html_content
-        # Verify that the query results are embedded into the HTML, allowing them to be visualized.
-        # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few
-        # sentinel strings within the query results, instead.
+        # Sanity check that query results are not embedded into the HTML.
         assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content
+            "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" not in html_content
         )  # identifier in 1st row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content
-        )  # identifier in 2nd row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
-        )  # identifier in 3rd row of query result
 
         # Make sure we can run a second graph query, after the graph server is already running.
         try:
@@ -704,18 +696,6 @@ def test_bigquery_graph_json_result(monkeypatch):
         html_content = display_mock.call_args_list[0][0][0].data
         assert "<script>" in html_content
         assert "</script>" in html_content
-        # Verify that the query results are embedded into the HTML, allowing them to be visualized.
-        # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few
-        # sentinel strings within the query results, instead.
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content
-        )  # identifier in 1st row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content
-        )  # identifier in 2nd row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
-        )  # identifier in 3rd row of query result
 
     assert bqstorage_mock.called  # BQ storage client was used
     assert isinstance(return_value, pandas.DataFrame)
@@ -791,18 +771,10 @@ def test_bigquery_graph_colab(monkeypatch):
         html_content = display_mock.call_args_list[0][0][0].data
         assert "<script>" in html_content
         assert "</script>" in html_content
-        # Verify that the query results are embedded into the HTML, allowing them to be visualized.
-        # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few
-        # sentinel strings within the query results, instead.
+        # Verify that the query results are not embedded into the HTML.
         assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content
+            "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" not in html_content
         )  # identifier in 1st row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content
-        )  # identifier in 2nd row of query result
-        assert (
-            "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
-        )  # identifier in 3rd row of query result
 
         # Make sure we actually used colab path, not GraphServer path.
         assert sys.modules["google.colab"].output.register_callback.called
@@ -818,6 +790,7 @@ def test_bigquery_graph_colab(monkeypatch):
     reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
 )
 def test_colab_query_callback():
+    graph_server.graph_server.query_result = pandas.DataFrame([], columns=["result"])
     result = bigquery_magics.bigquery._colab_query_callback(
         "query", json.dumps({"result": {}})
     )

tests/unit/test_graph_server.py

@@ -15,6 +15,7 @@
 import json
 import unittest
 
+import pandas as pd
 import pytest
 import requests
 
@@ -469,18 +470,14 @@ def test_post_ping(self):
     @pytest.mark.skipif(
         graph_visualization is None, reason="Requires `spanner-graph-notebook`"
     )
-    def test_post_query(self):
+    def test_post_query(self):        
         self.assertTrue(self.server_thread.is_alive())
+        graph_server.graph_server.query_result = pd.DataFrame([json.dumps(row_alex_owns_account)], columns=["result"])
         route = graph_server.graph_server.build_route(
             graph_server.GraphServer.endpoints["post_query"]
         )
 
-        data = {
-            "result": {
-                "0": json.dumps(row_alex_owns_account),
-            }
-        }
-        response = requests.post(route, json={"params": json.dumps(data)})
+        response = requests.post(route, json={"params": "{}"})
         self.assertEqual(response.status_code, 200)
         response_data = response.json()["response"]