fix: apply ruff formatting and remove unused imports in new test files

Author: mingjerli

tests/test_case_expressions.py

@@ -18,11 +18,8 @@
 Total: 16 test cases
 """
 
-import pytest
-
 from clgraph import RecursiveLineageBuilder, SQLColumnTracer
 
-
 # ============================================================================
 # Test Group 1: Simple CASE WHEN
 # ============================================================================
@@ -154,10 +151,10 @@ def test_case_references_multiple_tables(self):
         ref_edges = [e for e in graph.edges if e.to_node.full_name == "output.reference_code"]
         source_tables = {e.from_node.table_name for e in ref_edges}
 
-        assert "orders" in source_tables or "o" in source_tables, \
-            "Orders table should be a source"
-        assert "shipments" in source_tables or "s" in source_tables, \
+        assert "orders" in source_tables or "o" in source_tables, "Orders table should be a source"
+        assert "shipments" in source_tables or "s" in source_tables, (
             "Shipments table should be a source"
+        )
 
     def test_case_multi_table_backward_lineage(self):
         """Backward lineage should identify required inputs from all joined tables."""
@@ -182,8 +179,9 @@ def test_case_multi_table_backward_lineage(self):
             all_cols.update(cols)
 
         assert "role" in all_cols, "Condition column from users should be required"
-        assert "full_access_level" in all_cols or "basic_access_level" in all_cols, \
+        assert "full_access_level" in all_cols or "basic_access_level" in all_cols, (
             "Result columns from permissions should be required"
+        )
 
 
 # ============================================================================
@@ -294,8 +292,9 @@ def test_count_case_when(self):
         count_edges = [e for e in graph.edges if e.to_node.full_name == "output.active_count"]
         source_columns = {e.from_node.column_name for e in count_edges}
 
-        assert "is_active" in source_columns or "employee_id" in source_columns, \
+        assert "is_active" in source_columns or "employee_id" in source_columns, (
             "At least the condition or result column should be traced"
+        )
 
 
 # ============================================================================

tests/test_dialect_coverage.py

@@ -11,7 +11,7 @@
 
 import pytest
 
-from clgraph import Pipeline, RecursiveLineageBuilder, SQLColumnTracer
+from clgraph import Pipeline, RecursiveLineageBuilder
 
 # ============================================================================
 # Helper utilities
@@ -25,11 +25,7 @@ def _output_column_names(graph):
 
 def _source_columns_for(graph, output_col):
     """Return source column names that feed into *output_col*."""
-    return {
-        e.from_node.column_name
-        for e in graph.edges
-        if e.to_node.column_name == output_col
-    }
+    return {e.from_node.column_name for e in graph.edges if e.to_node.column_name == output_col}
 
 
 # ============================================================================
@@ -474,8 +470,7 @@ def test_cross_dialect_output_columns_match(self):
                 baseline_cols = business_cols
             else:
                 assert business_cols == baseline_cols, (
-                    f"Dialect {dialect} produced {business_cols}, "
-                    f"expected {baseline_cols}"
+                    f"Dialect {dialect} produced {business_cols}, expected {baseline_cols}"
                 )
 
 

tests/test_join_types.py

@@ -32,7 +32,6 @@
 
 from clgraph import RecursiveLineageBuilder, SQLColumnTracer
 
-
 # ============================================================================
 # Test Group 1: INNER JOIN
 # ============================================================================
@@ -56,10 +55,7 @@ def test_inner_join_columns_from_both_tables(self):
         assert "output.order_id" in graph.nodes
         assert "output.amount" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.id", "output.id") in edges_dict
         assert ("users.name", "output.name") in edges_dict
         assert ("orders.order_id", "output.order_id") in edges_dict
@@ -103,10 +99,7 @@ def test_left_join_nullable_side_columns_tracked(self):
         assert "output.name" in graph.nodes
         assert "output.bio" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.id", "output.id") in edges_dict
         assert ("profiles.bio", "output.bio") in edges_dict
 
@@ -145,10 +138,7 @@ def test_right_join_columns_tracked(self):
         assert "output.order_id" in graph.nodes
         assert "output.amount" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.name", "output.name") in edges_dict
         assert ("orders.order_id", "output.order_id") in edges_dict
 
@@ -188,10 +178,7 @@ def test_full_outer_join_both_sides_tracked(self):
         assert "output.order_id" in graph.nodes
         assert "output.amount" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.id", "output.id") in edges_dict
         assert ("orders.amount", "output.amount") in edges_dict
 
@@ -232,10 +219,7 @@ def test_cross_join_no_on_clause(self):
         assert "output.name" in graph.nodes
         assert "output.color_name" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.name", "output.name") in edges_dict
         assert ("colors.color_name", "output.color_name") in edges_dict
 
@@ -262,12 +246,8 @@ def test_self_join_with_aliases(self):
         assert "output.manager" in graph.nodes
 
         # Both output columns trace back to employees.name
-        employee_edges = [
-            e for e in graph.edges if e.to_node.full_name == "output.employee"
-        ]
-        manager_edges = [
-            e for e in graph.edges if e.to_node.full_name == "output.manager"
-        ]
+        employee_edges = [e for e in graph.edges if e.to_node.full_name == "output.employee"]
+        manager_edges = [e for e in graph.edges if e.to_node.full_name == "output.manager"]
 
         assert len(employee_edges) > 0
         assert len(manager_edges) > 0
@@ -311,10 +291,7 @@ def test_three_table_join_chain(self):
         assert "output.order_id" in graph.nodes
         assert "output.product_name" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.name", "output.name") in edges_dict
         assert ("orders.order_id", "output.order_id") in edges_dict
         assert ("products.product_name", "output.product_name") in edges_dict
@@ -362,9 +339,7 @@ def test_join_against_subquery(self):
         assert "output.total_amount" in graph.nodes
 
         # total_amount should trace back through the subquery to orders.amount
-        total_edges = [
-            e for e in graph.edges if e.to_node.full_name == "output.total_amount"
-        ]
+        total_edges = [e for e in graph.edges if e.to_node.full_name == "output.total_amount"]
         assert len(total_edges) > 0
 
     def test_join_subquery_backward_lineage(self):
@@ -406,12 +381,7 @@ def test_on_clause_columns_in_graph(self):
         assert "output.name" in graph.nodes
         assert "output.amount" in graph.nodes
 
-        # The join-key columns should appear as source nodes in the graph
-        all_source_names = {
-            e.from_node.full_name for e in graph.edges
-        }
-        # users.id and orders.user_id are join keys; at minimum the selected
-        # columns (users.name, orders.amount) must have edges
+        # At minimum the selected columns (users.name, orders.amount) must have edges
         assert ("users.name", "output.name") in {
             (e.from_node.full_name, e.to_node.full_name) for e in graph.edges
         }
@@ -447,10 +417,7 @@ def test_inner_join_lineage_per_dialect(self, dialect):
         assert "output.order_id" in graph.nodes
         assert "output.amount" in graph.nodes
 
-        edges_dict = {
-            (e.from_node.full_name, e.to_node.full_name): e
-            for e in graph.edges
-        }
+        edges_dict = {(e.from_node.full_name, e.to_node.full_name): e for e in graph.edges}
         assert ("users.id", "output.id") in edges_dict
         assert ("orders.amount", "output.amount") in edges_dict
 

tests/test_regex_functions.py

@@ -14,8 +14,7 @@
 
 import pytest
 
-from clgraph import RecursiveLineageBuilder, SQLColumnTracer
-
+from clgraph import RecursiveLineageBuilder
 
 # ============================================================================
 # Test Group 1: REGEXP_CONTAINS (BigQuery)
@@ -36,9 +35,7 @@ def test_regexp_contains_in_select(self):
         builder = RecursiveLineageBuilder(sql, dialect="bigquery")
         graph = builder.build()
 
-        is_example_edges = [
-            e for e in graph.edges if e.to_node.column_name == "is_example_domain"
-        ]
+        is_example_edges = [e for e in graph.edges if e.to_node.column_name == "is_example_domain"]
         assert len(is_example_edges) > 0
 
         source_columns = {e.from_node.column_name for e in is_example_edges}
@@ -81,9 +78,7 @@ def test_snowflake_regexp_substr(self):
         builder = RecursiveLineageBuilder(sql, dialect="snowflake")
         graph = builder.build()
 
-        area_code_edges = [
-            e for e in graph.edges if e.to_node.column_name == "area_code"
-        ]
+        area_code_edges = [e for e in graph.edges if e.to_node.column_name == "area_code"]
         assert len(area_code_edges) > 0
 
         source_columns = {e.from_node.column_name for e in area_code_edges}
@@ -108,9 +103,7 @@ def test_bigquery_regexp_replace(self):
         builder = RecursiveLineageBuilder(sql, dialect="bigquery")
         graph = builder.build()
 
-        clean_phone_edges = [
-            e for e in graph.edges if e.to_node.column_name == "clean_phone"
-        ]
+        clean_phone_edges = [e for e in graph.edges if e.to_node.column_name == "clean_phone"]
         assert len(clean_phone_edges) > 0
 
         source_columns = {e.from_node.column_name for e in clean_phone_edges}
@@ -126,9 +119,7 @@ def test_snowflake_regexp_replace(self):
         builder = RecursiveLineageBuilder(sql, dialect="snowflake")
         graph = builder.build()
 
-        sanitized_edges = [
-            e for e in graph.edges if e.to_node.column_name == "sanitized_address"
-        ]
+        sanitized_edges = [e for e in graph.edges if e.to_node.column_name == "sanitized_address"]
         assert len(sanitized_edges) > 0
 
         source_columns = {e.from_node.column_name for e in sanitized_edges}
@@ -178,9 +169,7 @@ def test_tilde_operator_in_where(self):
         builder = RecursiveLineageBuilder(sql, dialect="postgres")
         graph = builder.build()
 
-        desc_edges = [
-            e for e in graph.edges if e.to_node.column_name == "description"
-        ]
+        desc_edges = [e for e in graph.edges if e.to_node.column_name == "description"]
         assert len(desc_edges) > 0
 
         source_columns = {e.from_node.column_name for e in desc_edges}
@@ -209,9 +198,7 @@ def test_regexp_extract_in_case_when(self):
         builder = RecursiveLineageBuilder(sql, dialect="bigquery")
         graph = builder.build()
 
-        protocol_edges = [
-            e for e in graph.edges if e.to_node.column_name == "protocol_type"
-        ]
+        protocol_edges = [e for e in graph.edges if e.to_node.column_name == "protocol_type"]
         assert len(protocol_edges) > 0
 
         source_columns = {e.from_node.column_name for e in protocol_edges}
@@ -231,9 +218,7 @@ def test_regexp_replace_in_case_value(self):
         builder = RecursiveLineageBuilder(sql, dialect="bigquery")
         graph = builder.build()
 
-        phone_edges = [
-            e for e in graph.edges if e.to_node.column_name == "normalized_phone"
-        ]
+        phone_edges = [e for e in graph.edges if e.to_node.column_name == "normalized_phone"]
         assert len(phone_edges) > 0
 
         source_columns = {e.from_node.column_name for e in phone_edges}
@@ -262,15 +247,11 @@ def test_regexp_where_preserves_select_lineage(self):
         graph = builder.build()
 
         # user_id should still trace to source
-        user_id_edges = [
-            e for e in graph.edges if e.to_node.column_name == "user_id"
-        ]
+        user_id_edges = [e for e in graph.edges if e.to_node.column_name == "user_id"]
         assert len(user_id_edges) > 0
 
         # email_domain should trace back to email
-        domain_edges = [
-            e for e in graph.edges if e.to_node.column_name == "email_domain"
-        ]
+        domain_edges = [e for e in graph.edges if e.to_node.column_name == "email_domain"]
         assert len(domain_edges) > 0
 
         source_columns = {e.from_node.column_name for e in domain_edges}
@@ -295,9 +276,7 @@ def test_regexp_replace_bigquery(self):
         builder = RecursiveLineageBuilder(sql, dialect="bigquery")
         graph = builder.build()
 
-        clean_edges = [
-            e for e in graph.edges if e.to_node.column_name == "clean_name"
-        ]
+        clean_edges = [e for e in graph.edges if e.to_node.column_name == "clean_name"]
         assert len(clean_edges) > 0
 
         source_columns = {e.from_node.column_name for e in clean_edges}
@@ -313,9 +292,7 @@ def test_regexp_replace_postgres(self):
         builder = RecursiveLineageBuilder(sql, dialect="postgres")
         graph = builder.build()
 
-        clean_edges = [
-            e for e in graph.edges if e.to_node.column_name == "clean_name"
-        ]
+        clean_edges = [e for e in graph.edges if e.to_node.column_name == "clean_name"]
         assert len(clean_edges) > 0
 
         source_columns = {e.from_node.column_name for e in clean_edges}
@@ -331,9 +308,7 @@ def test_regexp_replace_snowflake(self):
         builder = RecursiveLineageBuilder(sql, dialect="snowflake")
         graph = builder.build()
 
-        clean_edges = [
-            e for e in graph.edges if e.to_node.column_name == "clean_name"
-        ]
+        clean_edges = [e for e in graph.edges if e.to_node.column_name == "clean_name"]
         assert len(clean_edges) > 0
 
         source_columns = {e.from_node.column_name for e in clean_edges}

tests/test_type_casting.py

@@ -267,9 +267,9 @@ def test_cast_lineage_across_dialects(self, dialect):
 
         edges = [e for e in graph.edges if e.to_node.column_name == "amount_int"]
         assert len(edges) > 0, f"No edges for amount_int in dialect={dialect}"
-        assert any(
-            e.from_node.column_name == "amount" for e in edges
-        ), f"Source column 'amount' not found in dialect={dialect}"
+        assert any(e.from_node.column_name == "amount" for e in edges), (
+            f"Source column 'amount' not found in dialect={dialect}"
+        )
 
     @pytest.mark.parametrize("dialect", ["bigquery", "postgres", "snowflake"])
     def test_cast_backward_lineage_across_dialects(self, dialect):