From c388b50d5e48b7b4bf60642bad83e002905a8640 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Thu, 30 Jan 2025 03:36:45 +0100 Subject: [PATCH] Remove import from MySQL provider tests in generic transfer test --- .../standard/operators/test_generic_transfer.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/providers/standard/tests/provider_tests/standard/operators/test_generic_transfer.py b/providers/standard/tests/provider_tests/standard/operators/test_generic_transfer.py index 4ea08e48891e6..3bdaa28def0e2 100644 --- a/providers/standard/tests/provider_tests/standard/operators/test_generic_transfer.py +++ b/providers/standard/tests/provider_tests/standard/operators/test_generic_transfer.py @@ -26,6 +26,7 @@ from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models.dag import DAG +from airflow.providers.mysql.hooks.mysql import MySqlHook from airflow.providers.postgres.hooks.postgres import PostgresHook from airflow.utils import timezone @@ -65,7 +66,17 @@ def teardown_method(self): ], ) def test_mysql_to_mysql(self, client): - from providers.tests.mysql.hooks.test_mysql import MySqlContext + class MySqlContext: + def __init__(self, client): + self.client = client + self.connection = MySqlHook.get_connection(MySqlHook.default_conn_name) + self.init_client = self.connection.extra_dejson.get("client", "mysqlclient") + + def __enter__(self): + self.connection.set_extra(f'{{"client": "{self.client}"}}') + + def __exit__(self, exc_type, exc_val, exc_tb): + self.connection.set_extra(f'{{"client": "{self.init_client}"}}') with MySqlContext(client): sql = "SELECT * FROM connection;"