chore(ingest/teradata): fix lint errors (#14055)

This commit is contained in:
Harshal Sheth 2025-07-11 16:01:10 -04:00 committed by GitHub
parent 9d89003604
commit 84ad365b37
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 78 additions and 42 deletions

View File

@ -315,11 +315,12 @@ class TestConfigurationIntegration:
source = TeradataSource(config, PipelineContext(run_id="test")) source = TeradataSource(config, PipelineContext(run_id="test"))
# Test database filtering logic # Test database filtering logic
with patch( with (
"datahub.ingestion.source.sql.teradata.create_engine" patch(
) as mock_create_engine, patch( "datahub.ingestion.source.sql.teradata.create_engine"
"datahub.ingestion.source.sql.teradata.inspect" ) as mock_create_engine,
) as mock_inspect: patch("datahub.ingestion.source.sql.teradata.inspect") as mock_inspect,
):
mock_engine = MagicMock() mock_engine = MagicMock()
mock_create_engine.return_value = mock_engine mock_create_engine.return_value = mock_engine
@ -464,10 +465,15 @@ class TestComplexQueryScenarios:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback", return_value=mock_result ),
patch.object(
source,
"_execute_with_cursor_fallback",
return_value=mock_result,
),
): ):
entries = list(source._fetch_lineage_entries_chunked()) entries = list(source._fetch_lineage_entries_chunked())
@ -640,11 +646,14 @@ class TestResourceManagement:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback" ),
) as mock_execute: patch.object(
source, "_execute_with_cursor_fallback"
) as mock_execute,
):
# Raise exception during query execution # Raise exception during query execution
mock_execute.side_effect = Exception("Database error") mock_execute.side_effect = Exception("Database error")

View File

@ -251,10 +251,15 @@ class TestMemoryOptimizations:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback", return_value=mock_result ),
patch.object(
source,
"_execute_with_cursor_fallback",
return_value=mock_result,
),
): ):
# Process entries # Process entries
entries = list(source._fetch_lineage_entries_chunked()) entries = list(source._fetch_lineage_entries_chunked())

View File

@ -489,11 +489,12 @@ class TestStageTracking:
config = TeradataConfig.parse_obj(_base_config()) config = TeradataConfig.parse_obj(_base_config())
# Create source without mocking to test the actual stage tracking during init # Create source without mocking to test the actual stage tracking during init
with patch( with (
"datahub.sql_parsing.sql_parsing_aggregator.SqlParsingAggregator" patch("datahub.sql_parsing.sql_parsing_aggregator.SqlParsingAggregator"),
), patch( patch(
"datahub.ingestion.source.sql.teradata.TeradataSource.cache_tables_and_views" "datahub.ingestion.source.sql.teradata.TeradataSource.cache_tables_and_views"
) as mock_cache: ) as mock_cache,
):
TeradataSource(config, PipelineContext(run_id="test")) TeradataSource(config, PipelineContext(run_id="test"))
# Verify cache_tables_and_views was called during init (stage tracking happens there) # Verify cache_tables_and_views was called during init (stage tracking happens there)
@ -775,11 +776,14 @@ class TestLineageQuerySeparation:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback" ),
) as mock_execute: patch.object(
source, "_execute_with_cursor_fallback"
) as mock_execute,
):
mock_execute.side_effect = [mock_result1, mock_result2] mock_execute.side_effect = [mock_result1, mock_result2]
entries = list(source._fetch_lineage_entries_chunked()) entries = list(source._fetch_lineage_entries_chunked())
@ -826,11 +830,16 @@ class TestLineageQuerySeparation:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback", return_value=mock_result ),
) as mock_execute: patch.object(
source,
"_execute_with_cursor_fallback",
return_value=mock_result,
) as mock_execute,
):
entries = list(source._fetch_lineage_entries_chunked()) entries = list(source._fetch_lineage_entries_chunked())
# Should have executed only one query # Should have executed only one query
@ -878,10 +887,15 @@ class TestLineageQuerySeparation:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback", return_value=mock_result ),
patch.object(
source,
"_execute_with_cursor_fallback",
return_value=mock_result,
),
): ):
entries = list(source._fetch_lineage_entries_chunked()) entries = list(source._fetch_lineage_entries_chunked())
@ -1005,11 +1019,19 @@ class TestLineageQuerySeparation:
mock_connection mock_connection
) )
with patch.object( with (
source, "get_metadata_engine", return_value=mock_engine patch.object(
), patch.object( source, "get_metadata_engine", return_value=mock_engine
source, "_execute_with_cursor_fallback", return_value=mock_result ),
), patch("datahub.ingestion.source.sql.teradata.logger") as mock_logger: 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()) list(source._fetch_lineage_entries_chunked())
# Verify progress logging for multiple queries # Verify progress logging for multiple queries