diff --git a/tests/sources/test_generic_database.py b/tests/sources/test_generic_database.py index 6103f9bcf..be47409bc 100644 --- a/tests/sources/test_generic_database.py +++ b/tests/sources/test_generic_database.py @@ -5,7 +5,9 @@ # """Tests the Generic Database source class methods""" +from contextlib import contextmanager from functools import partial +from unittest.mock import patch import pytest @@ -120,6 +122,15 @@ def fetchmany(self, size): return [] +@contextmanager +def mock_sync_db_engine(create_engine_target, query_object): + """Mock SQLAlchemy engine creation for sync database connector tests.""" + with patch(create_engine_target) as mock_create_engine: + mock_engine = mock_create_engine.return_value + mock_engine.connect.return_value = ConnectionSync(query_object) + yield mock_engine + + @pytest.mark.parametrize( "tables, expected_tables", [ diff --git a/tests/sources/test_graphql.py b/tests/sources/test_graphql.py index 75e61210f..611c9d73a 100644 --- a/tests/sources/test_graphql.py +++ b/tests/sources/test_graphql.py @@ -45,16 +45,22 @@ async def create_graphql_source( graphql_query="{users {name {firstName } } }", graphql_object_to_id_map='{"users": "id"}', ): - async with create_source( - GraphQLDataSource, - http_endpoint="http://127.0.0.1:1234", - authentication_method="none", - graphql_query=graphql_query, - graphql_object_to_id_map=graphql_object_to_id_map, - headers=headers, - graphql_variables=graphql_variables, - ) as source: - yield source + mock_session = AsyncMock() + mock_session.close = AsyncMock() + with patch( + "connectors.sources.graphql.aiohttp.ClientSession", + return_value=mock_session, + ): + async with create_source( + GraphQLDataSource, + http_endpoint="http://127.0.0.1:1234", + authentication_method="none", + graphql_query=graphql_query, + graphql_object_to_id_map=graphql_object_to_id_map, + headers=headers, + graphql_variables=graphql_variables, + ) as source: + yield source @pytest.mark.asyncio @@ -318,6 +324,16 @@ async def test_fetch_data_without_pageinfo(): pass +@pytest.mark.asyncio +async def test_teardown_does_not_use_real_client_session(): + """Regression: source teardown must not instantiate a real aiohttp.ClientSession.""" + with patch("aiohttp.ClientSession") as real_client_session: + async with create_graphql_source(): + pass + + real_client_session.assert_not_called() + + @pytest.mark.asyncio @freeze_time("2024-01-24T04:07:19") async def test_get_docs(): diff --git a/tests/sources/test_mssql.py b/tests/sources/test_mssql.py index 57a3dfc94..e7fe0cf73 100644 --- a/tests/sources/test_mssql.py +++ b/tests/sources/test_mssql.py @@ -9,7 +9,6 @@ from unittest.mock import ANY, AsyncMock, Mock, patch import pytest -from sqlalchemy.engine import Engine from sqlalchemy.exc import ProgrammingError from connectors.filtering.validation import SyncRuleValidationResult @@ -20,7 +19,9 @@ MSSQLQueries, ) from tests.sources.support import create_source -from tests.sources.test_generic_database import ConnectionSync +from tests.sources.test_generic_database import ConnectionSync, mock_sync_db_engine + +MSSQL_CREATE_ENGINE = "connectors.sources.mssql.create_engine" ADVANCED_SNIPPET = "advanced_snippet" @@ -44,10 +45,7 @@ def connect(self): @pytest.mark.asyncio async def test_ping(): async with create_source(MSSQLDataSource) as source: - source.engine = MockEngine() - with patch.object( - Engine, "connect", return_value=ConnectionSync(MSSQLQueries()) - ): + with mock_sync_db_engine(MSSQL_CREATE_ENGINE, MSSQLQueries()): await source.ping() @@ -56,7 +54,10 @@ async def test_ping(): async def test_ping_negative(): with pytest.raises(Exception): async with create_source(MSSQLDataSource) as source: - with patch.object(Engine, "connect", side_effect=Exception()): + with mock_sync_db_engine( + MSSQL_CREATE_ENGINE, MSSQLQueries() + ) as mock_engine: + mock_engine.connect.side_effect = Exception() await source.ping() @@ -151,9 +152,7 @@ async def test_get_docs(): async with create_source( MSSQLDataSource, database="xe", tables="*", schema="dbo" ) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(MSSQLQueries()) - ): + with mock_sync_db_engine(MSSQL_CREATE_ENGINE, MSSQLQueries()): actual_response = [] expected_response = [ { @@ -271,9 +270,7 @@ async def test_advanced_rules_validation(advanced_rules, expected_validation_res async with create_source( MSSQLDataSource, database="xe", tables="*", schema="dbo" ) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(MSSQLQueries()) - ): + with mock_sync_db_engine(MSSQL_CREATE_ENGINE, MSSQLQueries()): validation_result = await MSSQLAdvancedRulesValidator(source).validate( advanced_rules ) @@ -413,9 +410,7 @@ async def test_advanced_rules_validation_when_id_in_source_available( async with create_source( MSSQLDataSource, database="xe", tables="*", schema="dbo" ) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(MSSQLQueries()) - ): + with mock_sync_db_engine(MSSQL_CREATE_ENGINE, MSSQLQueries()): validation_result = await MSSQLAdvancedRulesValidator(source).validate( advanced_rules ) @@ -560,9 +555,7 @@ async def test_get_docs_with_advanced_rules(filtering, expected_response): async with create_source( MSSQLDataSource, database="xe", tables="*", schema="dbo" ) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(MSSQLQueries()) - ): + with mock_sync_db_engine(MSSQL_CREATE_ENGINE, MSSQLQueries()): actual_response = [] async for doc in source.get_docs(filtering=filtering): diff --git a/tests/sources/test_oracle.py b/tests/sources/test_oracle.py index bc411583d..793c86948 100644 --- a/tests/sources/test_oracle.py +++ b/tests/sources/test_oracle.py @@ -9,16 +9,16 @@ from unittest.mock import patch import pytest -from sqlalchemy.engine import Engine from connectors.sources.oracle import OracleClient, OracleDataSource, OracleQueries from tests.sources.support import create_source -from tests.sources.test_generic_database import ConnectionSync +from tests.sources.test_generic_database import mock_sync_db_engine DSN_SID = "oracle+oracledb://admin:Password_123@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST=127.0.0.1)(PORT=9090))(CONNECT_DATA=(SID=xe)))" DSN_SERVICE_NAME = "oracle+oracledb://admin:Password_123@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST=127.0.0.1)(PORT=9090))(CONNECT_DATA=(service_name=xe)))" SID = "sid" SERVICE_NAME = "service_name" +ORACLE_CREATE_ENGINE = "connectors.sources.oracle.create_engine" @contextmanager @@ -45,7 +45,7 @@ def oracle_client(**extras): client.close() -@patch("connectors.sources.oracle.create_engine") +@patch(ORACLE_CREATE_ENGINE) @pytest.mark.parametrize( "connection_source, DSN", [ @@ -65,7 +65,7 @@ def test_engine_in_thin_mode(mock_fun, connection_source, DSN): mock_fun.assert_called_with(DSN) -@patch("connectors.sources.oracle.create_engine") +@patch(ORACLE_CREATE_ENGINE) @pytest.mark.parametrize( "connection_source, DSN", [ @@ -93,9 +93,7 @@ def test_engine_in_thick_mode(mock_fun, connection_source, DSN): @pytest.mark.asyncio async def test_ping(): async with create_source(OracleDataSource) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(OracleQueries()) - ): + with mock_sync_db_engine(ORACLE_CREATE_ENGINE, OracleQueries()): await source.ping() @@ -110,9 +108,7 @@ async def test_get_docs(): sid="xe", tables="*", ) as source: - with patch.object( - Engine, "connect", return_value=ConnectionSync(OracleQueries()) - ): + with mock_sync_db_engine(ORACLE_CREATE_ENGINE, OracleQueries()): actual_response = [] expected_response = [ {