From 84ad365b3734bddac30e11e0b2ebbcbb871305ce Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Fri, 11 Jul 2025 16:01:10 -0400 Subject: [PATCH] chore(ingest/teradata): fix lint errors (#14055) --- .../tests/unit/test_teradata_integration.py | 37 ++++++---- .../tests/unit/test_teradata_performance.py | 13 ++-- .../tests/unit/test_teradata_source.py | 70 ++++++++++++------- 3 files changed, 78 insertions(+), 42 deletions(-) diff --git a/metadata-ingestion/tests/unit/test_teradata_integration.py b/metadata-ingestion/tests/unit/test_teradata_integration.py index 6ae7248cf3..710245e268 100644 --- a/metadata-ingestion/tests/unit/test_teradata_integration.py +++ b/metadata-ingestion/tests/unit/test_teradata_integration.py @@ -315,11 +315,12 @@ class TestConfigurationIntegration: source = TeradataSource(config, PipelineContext(run_id="test")) # Test database filtering logic - with patch( - "datahub.ingestion.source.sql.teradata.create_engine" - ) as mock_create_engine, patch( - "datahub.ingestion.source.sql.teradata.inspect" - ) as mock_inspect: + with ( + patch( + "datahub.ingestion.source.sql.teradata.create_engine" + ) as mock_create_engine, + patch("datahub.ingestion.source.sql.teradata.inspect") as mock_inspect, + ): mock_engine = MagicMock() mock_create_engine.return_value = mock_engine @@ -464,10 +465,15 @@ class TestComplexQueryScenarios: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback", return_value=mock_result + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, + "_execute_with_cursor_fallback", + return_value=mock_result, + ), ): entries = list(source._fetch_lineage_entries_chunked()) @@ -640,11 +646,14 @@ class TestResourceManagement: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback" - ) as mock_execute: + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, "_execute_with_cursor_fallback" + ) as mock_execute, + ): # Raise exception during query execution mock_execute.side_effect = Exception("Database error") diff --git a/metadata-ingestion/tests/unit/test_teradata_performance.py b/metadata-ingestion/tests/unit/test_teradata_performance.py index a3a91e60f6..738c20d108 100644 --- a/metadata-ingestion/tests/unit/test_teradata_performance.py +++ b/metadata-ingestion/tests/unit/test_teradata_performance.py @@ -251,10 +251,15 @@ class TestMemoryOptimizations: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback", return_value=mock_result + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, + "_execute_with_cursor_fallback", + return_value=mock_result, + ), ): # Process entries entries = list(source._fetch_lineage_entries_chunked()) diff --git a/metadata-ingestion/tests/unit/test_teradata_source.py b/metadata-ingestion/tests/unit/test_teradata_source.py index 4f29b625a5..8d440e6e96 100644 --- a/metadata-ingestion/tests/unit/test_teradata_source.py +++ b/metadata-ingestion/tests/unit/test_teradata_source.py @@ -489,11 +489,12 @@ class TestStageTracking: config = TeradataConfig.parse_obj(_base_config()) # Create source without mocking to test the actual stage tracking during init - with patch( - "datahub.sql_parsing.sql_parsing_aggregator.SqlParsingAggregator" - ), patch( - "datahub.ingestion.source.sql.teradata.TeradataSource.cache_tables_and_views" - ) as mock_cache: + with ( + patch("datahub.sql_parsing.sql_parsing_aggregator.SqlParsingAggregator"), + patch( + "datahub.ingestion.source.sql.teradata.TeradataSource.cache_tables_and_views" + ) as mock_cache, + ): TeradataSource(config, PipelineContext(run_id="test")) # Verify cache_tables_and_views was called during init (stage tracking happens there) @@ -775,11 +776,14 @@ class TestLineageQuerySeparation: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback" - ) as mock_execute: + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, "_execute_with_cursor_fallback" + ) as mock_execute, + ): mock_execute.side_effect = [mock_result1, mock_result2] entries = list(source._fetch_lineage_entries_chunked()) @@ -826,11 +830,16 @@ class TestLineageQuerySeparation: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback", return_value=mock_result - ) as mock_execute: + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, + "_execute_with_cursor_fallback", + return_value=mock_result, + ) as mock_execute, + ): entries = list(source._fetch_lineage_entries_chunked()) # Should have executed only one query @@ -878,10 +887,15 @@ class TestLineageQuerySeparation: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback", return_value=mock_result + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, + "_execute_with_cursor_fallback", + return_value=mock_result, + ), ): entries = list(source._fetch_lineage_entries_chunked()) @@ -1005,11 +1019,19 @@ class TestLineageQuerySeparation: mock_connection ) - with patch.object( - source, "get_metadata_engine", return_value=mock_engine - ), patch.object( - source, "_execute_with_cursor_fallback", return_value=mock_result - ), patch("datahub.ingestion.source.sql.teradata.logger") as mock_logger: + with ( + patch.object( + source, "get_metadata_engine", return_value=mock_engine + ), + patch.object( + source, + "_execute_with_cursor_fallback", + return_value=mock_result, + ), + patch( + "datahub.ingestion.source.sql.teradata.logger" + ) as mock_logger, + ): list(source._fetch_lineage_entries_chunked()) # Verify progress logging for multiple queries