Skip to content

Commit

Permalink
Remove import from MySQL provider tests in generic transfer test (apa…
Browse files Browse the repository at this point in the history
  • Loading branch information
potiuk authored and ambika-garg committed Jan 30, 2025
1 parent baac44d commit 4f4ed1d
Showing 1 changed file with 12 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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;"
Expand Down

0 comments on commit 4f4ed1d

Please sign in to comment.