Add more unit tests for convert_graph_data() to boost code coverage.

Author: ericfe-google

bigquery_magics/graph_server.py

@@ -64,18 +64,19 @@ def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
         for key, value in query_results.items():
             if column_name is None:
                 if not isinstance(key, str):
-                    raise ValueError(f"Expected key to be str, got {type(key)}")
+                    raise ValueError(f"Expected outer key to be str, got {type(key)}")
                 if not isinstance(value, dict):
-                    raise ValueError(f"Expected value to be dict, got {type(value)}")
+                    raise ValueError(f"Expected outer value to be dict, got {type(value)}")
                 column_name = key
                 column_value = value
             else:
+                # TODO: Implement multi-column support.
                 raise ValueError(
                     "Query has multiple columns - graph visualization not supported"
                 )
         if column_name is None or column_value is None:
             raise ValueError(
-                "Unable to get column name or value - how is this possible???"
+                "query result with no columns is not supported for graph visualization"
             )
 
         fields: List[StructType.Field] = [
@@ -85,9 +86,9 @@ def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
         rows = []
         for value_key, value_value in column_value.items():
             if not isinstance(value_key, str):
-                raise ValueError(f"Expected key to be str, got {type(key)}")
+                raise ValueError(f"Expected inner key to be str, got {type(value_key)}")
             if not isinstance(value_value, str):
-                raise ValueError(f"Expected value to be str, got {type(value)}")
+                raise ValueError(f"Expected inner value to be str, got {type(value_value)}")
             row_index = int(value_key)
             row_json = json.loads(value_value)
 

tests/unit/test_graph_server.py

@@ -148,7 +148,7 @@ def test_convert_one_column_no_rows():
 @pytest.mark.skipif(
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
-def test_convert_one_column_one_row():
+def test_convert_one_column_one_row_one_column():
     result = convert_graph_data(
         {
             "result": {
@@ -217,6 +217,82 @@ def test_convert_nongraph_json():
     assert result["response"]["schema"] is None
 
 
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_outer_key_not_string():
+    result = convert_graph_data(
+            {
+                0: {
+                    '0': json.dumps({"foo": 1, "bar": 2}),
+                }
+            })
+    assert result == {"error": "Expected outer key to be str, got <class 'int'>"}
+
+
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_outer_value_not_dict():
+    result = convert_graph_data(
+            {
+                'result': 0
+            })
+    assert result == {"error": "Expected outer value to be dict, got <class 'int'>"}
+
+
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_inner_key_not_string():
+    result = convert_graph_data(
+            {
+                'result': {
+                    0: json.dumps({"foo": 1, "bar": 2}),
+                }
+            })
+    assert result == {"error": "Expected inner key to be str, got <class 'int'>"}
+
+
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_inner_value_not_string():
+    result = convert_graph_data(
+            {
+                'result': {
+                    '0': 1,
+                }
+            })
+    assert result == {"error": "Expected inner value to be str, got <class 'int'>"}
+
+
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_one_column_one_row_two_columns():
+    result = convert_graph_data(
+        {
+            "result1": {
+                "0": json.dumps(row_alex_owns_account),
+            },
+            "result2": {
+                "0": json.dumps(row_alex_owns_account),
+            }
+
+        }
+    )
+    assert result == {"error": "Query has multiple columns - graph visualization not supported"}
+
+
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+def test_convert_empty_dict():
+    result = convert_graph_data({})
+    assert result == {"error": "query result with no columns is not supported for graph visualization"}
+
+
 class TestGraphServer(unittest.TestCase):
     def setUp(self):
         self.server_thread = GraphServer.init()