yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-11-02 11:49:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
metadata-ingestion/tests/unit/test_teradata_integration.py
View File

@ -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")

13
metadata-ingestion/tests/unit/test_teradata_performance.py
View File

@ -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())

70
metadata-ingestion/tests/unit/test_teradata_source.py
View File

@ -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