fix: resolve CI/CD issues - remove unused imports and fix type mismatches

tests/unit/test_graph_builder.py

@@ -1,7 +1,5 @@
-import pytest
 from pathlib import Path
 from src.graph.builder import GraphBuilder, GraphType, GraphNode, NodeType, GraphEdge
-from src.parsers.base import Entity, EntityType
 
 
 class TestGraphBuilder:
@@ -123,6 +121,19 @@ class TestGraphNode:
         assert node.style == "filled"
         assert node.shape == "ellipse"
 
+    def test_class_node_shape(self):
+        node = GraphNode(node_id="test", node_type=NodeType.CLASS, name="TestClass")
+        assert node.shape == "ellipse"
+
+    def test_class_node_shape_in_builder(self):
+        from src.graph.builder import GraphBuilder, GraphType
+        builder = GraphBuilder(GraphType.DIRECTED)
+        node = GraphNode(node_id="test", node_type=NodeType.CLASS, name="TestClass")
+        node.shape = "diamond"
+        builder.add_node(node)
+        added_node = builder.get_node_by_id("test")
+        assert added_node.shape == "diamond"
+
 
 class TestGraphEdge:
     def test_default_edge_type(self):