From 89837d45ec7582fe34e90d9fd7cec70e85bd842f Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 24 Jun 2025 14:07:52 +0530 Subject: [PATCH] Clean up messy default connection overrides in provider tests --- .../tests/unit/common/sql/hooks/test_sql.py | 20 +- .../unit/databricks/hooks/test_databricks.py | 413 ++++++++++-------- ...test_databricks_azure_workload_identity.py | 23 +- ...atabricks_azure_workload_identity_async.py | 23 +- .../databricks/hooks/test_databricks_sql.py | 20 +- .../databricks/triggers/test_databricks.py | 41 +- .../tests/unit/exasol/hooks/test_sql.py | 22 +- 7 files changed, 306 insertions(+), 256 deletions(-) diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py index 5d1aff2f1ed88..e881f52f69c45 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py @@ -33,7 +33,6 @@ from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.common.sql.hooks.sql import DbApiHook, resolve_dialects -from airflow.utils.session import provide_session from tests_common.test_utils.common_sql import mock_db_hook from tests_common.test_utils.providers import get_provider_min_airflow_version @@ -49,15 +48,18 @@ class DBApiHookForTests(DbApiHook): get_conn = MagicMock(name="conn") -@provide_session @pytest.fixture(autouse=True) -def create_connection(session): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.login = None - conn.password = PASSWORD - conn.extra = None - session.commit() +def create_connection(create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="sqlite", + host=HOST, + login=None, + password=PASSWORD, + extra=None, + ) + ) def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py index d3a85942ca8fb..2698a12b62946 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py @@ -53,7 +53,6 @@ TOKEN_REFRESH_LEAD_TIME, BearerAuth, ) -from airflow.utils.session import provide_session TASK_ID = "databricks-operator" DEFAULT_CONN_ID = "databricks_default" @@ -318,16 +317,17 @@ class TestDatabricksHook: Tests for DatabricksHook. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = LOGIN - conn.password = PASSWORD - conn.extra = None - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + ) + ) self.hook = DatabricksHook(retry_delay=0) @@ -1255,17 +1255,18 @@ class TestDatabricksHookToken: Tests for DatabricksHook when auth is done with token. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = None - conn.password = None - conn.extra = json.dumps({"token": TOKEN, "host": HOST}) - - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=None, + extra=json.dumps({"token": TOKEN, "host": HOST}), + ) + ) self.hook = DatabricksHook() @@ -1290,16 +1291,18 @@ class TestDatabricksHookTokenInPassword: Tests for DatabricksHook. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = None - conn.password = TOKEN - conn.extra = None - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=TOKEN, + extra=None, + ) + ) self.hook = DatabricksHook(retry_delay=0) @@ -1320,17 +1323,18 @@ def test_submit_run(self, mock_requests): @pytest.mark.db_test class TestDatabricksHookTokenWhenNoHostIsProvidedInExtra(TestDatabricksHookToken): - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = None - conn.password = None - conn.extra = json.dumps({"token": TOKEN}) - - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=None, + extra=json.dumps({"token": TOKEN}), + ) + ) self.hook = DatabricksHook() @@ -1341,16 +1345,20 @@ class TestDatabricksHookConnSettings(TestDatabricksHookToken): Tests that `schema` and/or `port` get reflected in the requested API URLs. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = "http" - conn.port = 7908 - conn.login = None - conn.password = None - conn.extra = json.dumps({"token": TOKEN}) - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=None, + extra=json.dumps({"token": TOKEN}), + schema="http", + port=7908, + ) + ) self.hook = DatabricksHook() @@ -1467,21 +1475,23 @@ class TestDatabricksHookAadToken: Tests for DatabricksHook when auth is done with AAD token for SP as user inside workspace. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = None - conn.schema = None - conn.port = None - conn.login = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn.password = "secret" - conn.extra = json.dumps( - { - "host": HOST, - "azure_tenant_id": "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d", - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=None, + login="9ff815a6-4404-4ab8-85cb-cd0e6f879c1d", + password="secret", + extra=json.dumps( + { + "azure_tenant_id": "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d", + } + ), + ) ) - session.commit() + self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") @@ -1508,25 +1518,26 @@ class TestDatabricksHookAadTokenOtherClouds: using non-global Azure cloud (China, GovCloud, Germany) """ - @provide_session - def setup_method(self, method, session=None): + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): self.tenant_id = "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d" self.ad_endpoint = "https://login.microsoftonline.de" self.client_id = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = None - conn.schema = None - conn.port = None - conn.login = self.client_id - conn.password = "secret" - conn.extra = json.dumps( - { - "host": HOST, - "azure_tenant_id": self.tenant_id, - "azure_ad_endpoint": self.ad_endpoint, - } + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=None, + login=self.client_id, + password="secret", + extra=json.dumps( + { + "azure_tenant_id": self.tenant_id, + "azure_ad_endpoint": self.ad_endpoint, + } + ), + ) ) - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") @@ -1560,23 +1571,25 @@ class TestDatabricksHookAadTokenSpOutside: Tests for DatabricksHook when auth is done with AAD token for SP outside of workspace. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): self.tenant_id = "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d" self.client_id = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn.login = self.client_id - conn.password = "secret" - conn.host = HOST - conn.schema = None - conn.port = None - conn.extra = json.dumps( - { - "azure_resource_id": "/Some/resource", - "azure_tenant_id": "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d", - } + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=self.client_id, + password="secret", + extra=json.dumps( + { + "azure_resource_id": "/Some/resource", + "azure_tenant_id": self.tenant_id, + } + ), + ) ) - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") @@ -1616,18 +1629,22 @@ class TestDatabricksHookAadTokenManagedIdentity: Tests for DatabricksHook when auth is done with AAD leveraging Managed Identity authentication """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.extra = json.dumps( - { - "use_azure_managed_identity": True, - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=None, + extra=json.dumps( + { + "use_azure_managed_identity": True, + } + ), + ) ) - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") @@ -1663,16 +1680,18 @@ class TestDatabricksHookAsyncMethods: Tests for async functionality of DatabricksHook. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = LOGIN - conn.password = PASSWORD - conn.extra = None - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + extra=None, + ) + ) self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @@ -1846,20 +1865,23 @@ class TestDatabricksHookAsyncAadToken: auth is done with AAD token for SP as user inside workspace. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.schema = None - conn.port = None - conn.login = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn.password = "secret" - conn.extra = json.dumps( - { - "host": HOST, - "azure_tenant_id": "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d", - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login="9ff815a6-4404-4ab8-85cb-cd0e6f879c1d", + password="secret", + extra=json.dumps( + { + "azure_tenant_id": "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d", + } + ), + ) ) - session.commit() + self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @pytest.mark.asyncio @@ -1889,24 +1911,26 @@ class TestDatabricksHookAsyncAadTokenOtherClouds: for SP as user inside workspace and using non-global Azure cloud (China, GovCloud, Germany) """ - @provide_session - def setup_method(self, method, session=None): + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): self.tenant_id = "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d" self.ad_endpoint = "https://login.microsoftonline.de" self.client_id = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.schema = None - conn.port = None - conn.login = self.client_id - conn.password = "secret" - conn.extra = json.dumps( - { - "host": HOST, - "azure_tenant_id": self.tenant_id, - "azure_ad_endpoint": self.ad_endpoint, - } + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=self.client_id, + password="secret", + extra=json.dumps( + { + "azure_tenant_id": self.tenant_id, + "azure_ad_endpoint": self.ad_endpoint, + } + ), + ) ) - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @pytest.mark.asyncio @@ -1945,23 +1969,25 @@ class TestDatabricksHookAsyncAadTokenSpOutside: Tests for DatabricksHook using async methods when auth is done with AAD token for SP outside of workspace. """ - @provide_session - def setup_method(self, method, session=None): + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): self.tenant_id = "3ff810a6-5504-4ab8-85cb-cd0e6f879c1d" self.client_id = "9ff815a6-4404-4ab8-85cb-cd0e6f879c1d" - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.login = self.client_id - conn.password = "secret" - conn.host = HOST - conn.schema = None - conn.port = None - conn.extra = json.dumps( - { - "azure_resource_id": "/Some/resource", - "azure_tenant_id": self.tenant_id, - } + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=self.client_id, + password="secret", + extra=json.dumps( + { + "azure_resource_id": "/Some/resource", + "azure_tenant_id": self.tenant_id, + } + ), + ) ) - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @pytest.mark.asyncio @@ -2008,19 +2034,22 @@ class TestDatabricksHookAsyncAadTokenManagedIdentity: auth is done with AAD leveraging Managed Identity authentication """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.extra = json.dumps( - { - "use_azure_managed_identity": True, - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=None, + extra=json.dumps( + { + "use_azure_managed_identity": True, + } + ), + ) ) - session.commit() - session.commit() self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @pytest.mark.asyncio @@ -2060,16 +2089,18 @@ class TestDatabricksHookSpToken: Tests for DatabricksHook when auth is done with Service Principal Oauth token. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = "c64f6d12-f6e4-45a4-846e-032b42b27758" - conn.password = "secret" - conn.extra = json.dumps({"service_principal_oauth": True}) - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login="c64f6d12-f6e4-45a4-846e-032b42b27758", + password="secret", + extra=json.dumps({"service_principal_oauth": True}), + ) + ) self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") @@ -2101,16 +2132,18 @@ class TestDatabricksHookAsyncSpToken: Principal Oauth token. """ - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.schema = None - conn.port = None - conn.login = "c64f6d12-f6e4-45a4-846e-032b42b27758" - conn.password = "secret" - conn.extra = json.dumps({"service_principal_oauth": True}) - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login="c64f6d12-f6e4-45a4-846e-032b42b27758", + password="secret", + extra=json.dumps({"service_principal_oauth": True}), + ) + ) self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) @pytest.mark.asyncio diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity.py index 6a57a1340cb42..d009fee7a6f4a 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity.py @@ -28,7 +28,6 @@ from airflow.models import Connection from airflow.providers.databricks.hooks.databricks import DatabricksHook from airflow.providers.databricks.hooks.databricks_base import DEFAULT_AZURE_CREDENTIAL_SETTING_KEY -from airflow.utils.session import provide_session def create_successful_response_mock(content): @@ -55,16 +54,20 @@ def create_aad_token_for_resource() -> AccessToken: class TestDatabricksHookAadTokenWorkloadIdentity: _hook: DatabricksHook - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.extra = json.dumps( - { - DEFAULT_AZURE_CREDENTIAL_SETTING_KEY: True, - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + extra=json.dumps( + { + DEFAULT_AZURE_CREDENTIAL_SETTING_KEY: True, + } + ), + ) ) - session.commit() # This will use the default connection id (databricks_default) self._hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity_async.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity_async.py index 0a43262acc4f6..fc297800aa14f 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity_async.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_azure_workload_identity_async.py @@ -28,7 +28,6 @@ from airflow.models import Connection from airflow.providers.databricks.hooks.databricks import DatabricksHook from airflow.providers.databricks.hooks.databricks_base import DEFAULT_AZURE_CREDENTIAL_SETTING_KEY -from airflow.utils.session import provide_session def create_successful_response_mock(content): @@ -53,16 +52,20 @@ def create_aad_token_for_resource() -> AccessToken: @pytest.mark.db_test class TestDatabricksHookAadTokenWorkloadIdentityAsync: - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.extra = json.dumps( - { - DEFAULT_AZURE_CREDENTIAL_SETTING_KEY: True, - } + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + extra=json.dumps( + { + DEFAULT_AZURE_CREDENTIAL_SETTING_KEY: True, + } + ), + ) ) - session.commit() @pytest.mark.asyncio @mock.patch( diff --git a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py index 4fd5d80c6521e..7449489c77d7b 100644 --- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py +++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py @@ -33,7 +33,6 @@ from airflow.models import Connection from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook, create_timeout_thread -from airflow.utils.session import provide_session TASK_ID = "databricks-sql-operator" DEFAULT_CONN_ID = "databricks_default" @@ -42,15 +41,18 @@ TOKEN = "token" -@provide_session @pytest.fixture(autouse=True) -def create_connection(session): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.login = None - conn.password = TOKEN - conn.extra = None - session.commit() +def create_connection(create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=None, + password=TOKEN, + extra=None, + ) + ) @pytest.fixture diff --git a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py index bb143fe5c9b0c..a34e4723388c1 100644 --- a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py @@ -29,7 +29,6 @@ DatabricksSQLStatementExecutionTrigger, ) from airflow.triggers.base import TriggerEvent -from airflow.utils.session import provide_session pytestmark = pytest.mark.db_test @@ -121,14 +120,18 @@ class TestDatabricksExecutionTrigger: - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.login = LOGIN - conn.password = PASSWORD - conn.extra = None - session.commit() + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + extra=None, + ) + ) self.trigger = DatabricksExecutionTrigger( run_id=RUN_ID, @@ -258,15 +261,19 @@ async def test_sleep_between_retries( class TestDatabricksSQLStatementExecutionTrigger: - @provide_session - def setup_method(self, method, session=None): + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): self.end_time = time.time() + 60 - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.login = LOGIN - conn.password = PASSWORD - conn.extra = None - session.commit() + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="databricks", + host=HOST, + login=LOGIN, + password=PASSWORD, + extra=None, + ) + ) self.trigger = DatabricksSQLStatementExecutionTrigger( statement_id=STATEMENT_ID, diff --git a/providers/exasol/tests/unit/exasol/hooks/test_sql.py b/providers/exasol/tests/unit/exasol/hooks/test_sql.py index da99da330e567..d8d8a4f37c2cc 100644 --- a/providers/exasol/tests/unit/exasol/hooks/test_sql.py +++ b/providers/exasol/tests/unit/exasol/hooks/test_sql.py @@ -27,7 +27,6 @@ from airflow.models import Connection from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.exasol.hooks.exasol import ExasolHook -from airflow.utils.session import provide_session TASK_ID = "sql-operator" HOST = "host" @@ -40,17 +39,18 @@ class ExasolHookForTests(ExasolHook): get_conn = MagicMock(name="conn") -@provide_session @pytest.fixture(autouse=True) -def create_connection(session): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - if conn is None: - conn = Connection(conn_id=DEFAULT_CONN_ID) - conn.host = HOST - conn.login = None - conn.password = PASSWORD - conn.extra = None - session.commit() +def create_connection(create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=DEFAULT_CONN_ID, + conn_type="exasol", + host=HOST, + login=None, + password=PASSWORD, + extra=None, + ) + ) @pytest.fixture