Fix broken tests and get build back to green

1. MagicsTest.test_spanner_graph_magic_with_empty_cell:
    - Added check for empty cell in magics.py to fix
    - Hung test fixed by adding tearDown method
2. ConversionTest.test_get_nodes_edges():
    - Incorrect use of DB.execute_query() API. Fixed

3. GraphEntitiesTest.test_add_edge_to_graph():
    - Wrong type in assertion. Changed from 'int' to 'str'

4. SampleNotebookTest.test_notebook_cells():
    - Wrong assertion (formatting)
    - Hung test fixed by adding tearDown method

5. GraphServerTest (multiple failures):
    - Formatting issue. Changed `'''` to `"` in execute_node_expansion()

Author: Sailesh Mukil Gangatharan

spanner_graphs/graph_server.py

@@ -56,27 +56,27 @@ class EdgeDirection(Enum):
 def validate_property_type(property_type: str) -> TypeCode:
     """
     Validates and converts a property type string to a Spanner TypeCode.
-    
+
     Args:
         property_type: The property type string from the request
-        
+
     Returns:
         The corresponding TypeCode enum value
-        
+
     Raises:
         ValueError: If the property type is invalid
     """
     if not property_type:
         raise ValueError("Property type must be provided")
-        
+
     # Convert to uppercase for case-insensitive comparison
     property_type = property_type.upper()
-    
+
     # Check if the type is valid
     if property_type not in PROPERTY_TYPE_MAP:
         valid_types = ', '.join(sorted(PROPERTY_TYPE_MAP.keys()))
         raise ValueError(f"Invalid property type: {property_type}. Allowed types are: {valid_types}")
-    
+
     return PROPERTY_TYPE_MAP[property_type]
 
 def validate_node_expansion_request(data) -> (list[NodePropertyForDataExploration], EdgeDirection):
@@ -149,11 +149,11 @@ def execute_node_expansion(
     params_str: str,
     request: dict) -> dict:
     """Execute a node expansion query to find connected nodes and edges.
-    
+
     Args:
         params_str: A JSON string containing connection parameters (project, instance, database, graph, mock).
         request: A dictionary containing node expansion request details (uid, node_labels, node_properties, direction, edge_label).
-    
+
     Returns:
         dict: A dictionary containing the query response with nodes and edges.
     """
@@ -188,7 +188,7 @@ def execute_node_expansion(
         if node_property.type in (TypeCode.INT64, TypeCode.NUMERIC, TypeCode.FLOAT32, TypeCode.FLOAT64, TypeCode.BOOL):
             value_str = node_property.value
         else:
-            value_str = f"\'''{node_property.value}\'''"
+            value_str = f'"{node_property.value}"'
         node_property_strings.append(f"n.{node_property.key}={value_str}")
 
     query = f"""
@@ -235,7 +235,7 @@ def execute_query(project: str, instance: str, database: str, query: str, mock =
                     "error": f"Query error: \n{getattr(err, 'message', str(err))}"
                 }
         nodes, edges = get_nodes_edges(query_result, fields, schema_json)
-        
+
         return {
             "response": {
                 "nodes": [node.to_json() for node in nodes],
@@ -360,7 +360,7 @@ def handle_post_ping(self):
     def handle_post_query(self):
         data = self.parse_post_data()
         params = json.loads(data["params"])
-        response = execute_query(            
+        response = execute_query(
             project=params["project"],
             instance=params["instance"],
             database=params["database"],
@@ -371,7 +371,7 @@ def handle_post_query(self):
 
     def handle_post_node_expansion(self):
         """Handle POST requests for node expansion.
-        
+
         Expects a JSON payload with:
         - params: A JSON string containing connection parameters (project, instance, database, graph)
         - request: A dictionary with node details (uid, node_labels, node_properties, direction, edge_label)

spanner_graphs/magics.py

@@ -94,7 +94,7 @@ def receive_query_request(query: str, params: str):
 
 def receive_node_expansion_request(request: dict, params_str: str):
     """Handle node expansion requests in Google Colab environment
-    
+
     Args:
         request: A dictionary containing node expansion details including:
             - uid: str - Unique identifier of the node to expand
@@ -108,7 +108,7 @@ def receive_node_expansion_request(request: dict, params_str: str):
             - database: str - Spanner database ID
             - graph: str - Graph name
             - mock: bool - Whether to use mock data
-    
+
     Returns:
         JSON: A JSON-serialized response containing either:
             - The query results with nodes and edges
@@ -165,6 +165,7 @@ def visualize(self):
     @cell_magic
     def spanner_graph(self, line: str, cell: str):
         """spanner_graph function"""
+
         parser = argparse.ArgumentParser(
             description="Visualize network from Spanner database",
             exit_on_error=False)
@@ -184,6 +185,9 @@ def spanner_graph(self, line: str, cell: str):
                     raise ValueError(
                         "Please provide `--project`, `--instance`, "
                         "and `--database` values for your query.")
+                if not cell or not cell.strip():
+                    print("Error: Query is required.")
+                    return
 
             self.args = parser.parse_args(line.split())
             self.cell = cell

tests/conversion_test.py

@@ -38,7 +38,7 @@ def test_get_nodes_edges(self) -> None:
         """
         # Get data from mock database
         mock_db = MockSpannerDatabase()
-        data, fields, _, schema_json = mock_db.execute_query("")
+        data, fields, _, schema_json, _ = mock_db.execute_query("")
 
         # Convert data to nodes and edges
         nodes, edges = get_nodes_edges(data, fields)
@@ -72,7 +72,7 @@ def test_get_nodes_edges(self) -> None:
             self.assertTrue(hasattr(edge, 'destination'), "Edge should have a destination")
             self.assertIsInstance(edge.labels, list, "Edge labels should be a list")
             self.assertIsInstance(edge.properties, dict, "Edge properties should be a dict")
-            
+
             # Verify edge endpoints exist in nodes
             source_exists = any(node.identifier == edge.source for node in nodes)
             dest_exists = any(node.identifier == edge.destination for node in nodes)
@@ -94,32 +94,32 @@ def test_get_nodes_edges_with_missing_nodes(self) -> None:
                 }),
                 json.dumps({
                     "kind": "node",
-                    "identifier": "node1", 
+                    "identifier": "node1",
                     "labels": ["Device"],
                     "properties": {"name": "Router"}
                 })
                 # Note: node2 is intentionally missing
             ]
         }
-        
+
         # Create a mock field for the column
         field = StructType.Field(
             name="column1",
             type_=Type(code=TypeCode.JSON)
         )
-        
+
         # Convert data to nodes and edges
         nodes, edges = get_nodes_edges(data, [field])
-        
+
         # Verify we got the expected number of nodes and edges
         self.assertEqual(len(edges), 1, "Should have one edge")
         self.assertEqual(len(nodes), 2, "Should have two nodes (one real, one intermediate)")
-        
+
         # Verify node identifiers
         node_ids = {node.identifier for node in nodes}
         self.assertIn("node1", node_ids, "Original node should exist")
         self.assertIn("node2", node_ids, "Missing node should be created as intermediate")
-        
+
         # Find the intermediate node
         intermediate_node = next((node for node in nodes if node.identifier == "node2"), None)
         self.assertIsNotNone(intermediate_node, "Intermediate node should exist")
@@ -150,33 +150,33 @@ def test_get_nodes_edges_with_multiple_references(self) -> None:
                 }),
                 json.dumps({
                     "kind": "node",
-                    "identifier": "node1", 
+                    "identifier": "node1",
                     "labels": ["Device"],
                     "properties": {"name": "Router"}
                 }),
                 json.dumps({
                     "kind": "node",
-                    "identifier": "node2", 
+                    "identifier": "node2",
                     "labels": ["Device"],
                     "properties": {"name": "Switch"}
                 })
                 # Note: missing_node is intentionally missing
             ]
         }
-        
+
         # Create a mock field for the column
         field = StructType.Field(
             name="column1",
             type_=Type(code=TypeCode.JSON)
         )
-        
+
         # Convert data to nodes and edges
         nodes, edges = get_nodes_edges(data, [field])
-        
+
         # Verify we got the expected number of nodes and edges
         self.assertEqual(len(edges), 2, "Should have two edges")
         self.assertEqual(len(nodes), 3, "Should have three nodes (two real, one intermediate)")
-        
+
         # Count intermediate nodes
         intermediate_nodes = [node for node in nodes if node.intermediate]
         self.assertEqual(len(intermediate_nodes), 1, "Should create only one intermediate node")
@@ -197,32 +197,32 @@ def test_get_nodes_edges_with_complete_data(self) -> None:
                 }),
                 json.dumps({
                     "kind": "node",
-                    "identifier": "node1", 
+                    "identifier": "node1",
                     "labels": ["Device"],
                     "properties": {"name": "Router"}
                 }),
                 json.dumps({
                     "kind": "node",
-                    "identifier": "node2", 
+                    "identifier": "node2",
                     "labels": ["Device"],
                     "properties": {"name": "Switch"}
                 })
             ]
         }
-        
+
         # Create a mock field for the column
         field = StructType.Field(
             name="column1",
             type_=Type(code=TypeCode.JSON)
         )
-        
+
         # Convert data to nodes and edges
         nodes, edges = get_nodes_edges(data, [field])
-        
+
         # Verify we got the expected number of nodes and edges
         self.assertEqual(len(edges), 1, "Should have one edge")
         self.assertEqual(len(nodes), 2, "Should have exactly two nodes (no intermediates)")
-        
+
         # Verify no intermediate nodes exist
         intermediate_nodes = [node for node in nodes if node.intermediate]
         self.assertEqual(len(intermediate_nodes), 0, "Should not create any intermediate nodes")

tests/graph_entities_test.py

@@ -43,11 +43,11 @@ def test_intermediate_flag_in_constructor(self):
         # Test with intermediate=True
         node1 = Node("1", ["Person"], {"name": "Emmanuel"}, intermediate=True)
         self.assertTrue(node1.intermediate, "Node should be marked as intermediate")
-        
+
         # Test with intermediate=False
         node2 = Node("2", ["Person"], {"name": "John"}, intermediate=False)
         self.assertFalse(node2.intermediate, "Node should not be marked as intermediate")
-        
+
         # Test with default (should be False)
         node3 = Node("3", ["Person"], {"name": "Alice"})
         self.assertFalse(node3.intermediate, "Node should default to not intermediate")
@@ -56,7 +56,7 @@ def test_make_intermediate(self):
         """Test the make_intermediate static method"""
         test_identifier = "test123"
         node = Node.make_intermediate(test_identifier)
-        
+
         self.assertEqual(node.identifier, test_identifier, "Identifier should match input")
         self.assertEqual(node.labels, ["Intermediate"], "Labels should include 'Intermediate'")
         self.assertTrue(node.intermediate, "Node should be marked as intermediate")
@@ -72,7 +72,7 @@ def test_to_json_with_intermediate(self):
         node1 = Node("1", ["Person"], {"name": "Jill"}, intermediate=True)
         json1 = node1.to_json()
         self.assertTrue(json1["intermediate"], "JSON should include intermediate=True")
-        
+
         # Test with intermediate=False
         node2 = Node("2", ["Person"], {"name": "John"}, intermediate=False)
         json2 = node2.to_json()
@@ -117,7 +117,7 @@ def test_add_edge_to_graph(self):
         edge = Edge.from_json(data)
         edge.add_to_graph(graph)
 
-        self.assertIn((1, 2), graph.edges)
+        self.assertIn(('1', '2'), graph.edges)
         # self.assertEqual(graph.edges[1, 2]["label"], "KNOWS")
         # self.assertEqual(graph.edges[1, 2]["title"],
         #                  "--- Edge Properties ---\nsince: 2020")