Importing correct MySqlContext in test_generic_transfer

providers/mysql/tests/provider_tests/mysql/hooks/test_mysql.py

@@ -29,6 +29,8 @@
 from airflow.models import Connection
 from airflow.models.dag import DAG
 
+from tests_common.test_utils.common_sql import MySqlContext
+
 try:
     import MySQLdb.cursors
 
@@ -412,19 +414,6 @@ def test_generate_insert_sql_with_already_escaped_column_name(self):
 TEST_DAG_ID = "unit_test_dag"
 
 
-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}"}}')
-
-
 @pytest.mark.backend("mysql")
 @pytest.mark.skipif(not MYSQL_AVAILABLE, reason="MySQL not available")
 class TestMySql:

providers/standard/tests/provider_tests/standard/operators/test_generic_transfer.py

@@ -65,7 +65,7 @@ def teardown_method(self):
         ],
     )
     def test_mysql_to_mysql(self, client):
-        from providers.tests.mysql.hooks.test_mysql import MySqlContext
+        from tests_common.test_utils.common_sql import MySqlContext
 
         with MySqlContext(client):
             sql = "SELECT * FROM connection;"

tests_common/test_utils/common_sql.py

@@ -21,6 +21,7 @@
 from unittest import mock
 
 from airflow.models import Connection
+from airflow.providers.mysql.hooks.mysql import MySqlHook
 
 if TYPE_CHECKING:
     from airflow.hooks.base import BaseHook
@@ -55,3 +56,18 @@ def get_conn(self):
             return conn
 
     return MockedHook(**hook_params)
+
+
+class MySqlContext:
+    """MySqlContext to be used in tests as context manager."""
+
+    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}"}}')