mirror of
https://github.com/datahub-project/datahub.git
synced 2025-11-17 20:03:59 +00:00
chore(ingest/teradata): fix lint errors (#14055)
This commit is contained in:
parent
9d89003604
commit
84ad365b37
@ -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")
|
||||||
|
|
||||||
|
|||||||
@ -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())
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user