Skip to content

Commit 5ea0d97

Browse files
XD-DENGFokko
authored andcommitted
[AIRFLOW-3239] Enable existing CI tests (#4131)
1. Renamed files: - tests/configuration.py → tests/test_configuration.py - tests/impersonation.py → tests/test_impersonation.py - tests/utils.py → tests/test_utils.py - tests/operators/operators.py → tests/operators/test_operators.py - tests/operators/bash_operator.py → tests/operators/test_bash_operator.py - tests/jobs.py → tests/test_jobs.py 2. Updated tests/__init__.py accordingly 3. Fixed database-specific tests in tests/operators/test_operators.py 4. Fixed issue in tests/operators/test_bash_operator.py
1 parent e53182c commit 5ea0d97

7 files changed

+51
-10
lines changed

tests/__init__.py

-10
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,5 @@
2222
from __future__ import absolute_import
2323

2424
from .api import *
25-
from .configuration import *
26-
from .contrib import *
2725
from .core import *
28-
from .executors import *
29-
from .jobs import *
30-
from .impersonation import *
31-
from .lineage import *
3226
from .models import *
33-
from .operators import *
34-
from .security import *
35-
from .task import *
36-
from .utils import *

tests/operators/bash_operator.py tests/operators/test_bash_operator.py

+5
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,9 @@ def test_echo_env_variables(self):
6565
'echo $AIRFLOW_CTX_EXECUTION_DATE>> {0};'
6666
'echo $AIRFLOW_CTX_DAG_RUN_ID>> {0};'.format(fname)
6767
)
68+
69+
original_AIRFLOW_HOME = os.environ['AIRFLOW_HOME']
70+
6871
os.environ['AIRFLOW_HOME'] = 'MY_PATH_TO_AIRFLOW_HOME'
6972
t.run(DEFAULT_DATE, DEFAULT_DATE,
7073
ignore_first_depends_on_past=True, ignore_ti_state=True)
@@ -78,3 +81,5 @@ def test_echo_env_variables(self):
7881
self.assertIn('echo_env_vars', output)
7982
self.assertIn(DEFAULT_DATE.isoformat(), output)
8083
self.assertIn('manual__' + DEFAULT_DATE.isoformat(), output)
84+
85+
os.environ['AIRFLOW_HOME'] = original_AIRFLOW_HOME

tests/operators/operators.py tests/operators/test_operators.py

+46
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ def tearDown(self):
5353
for table in drop_tables:
5454
conn.execute("DROP TABLE IF EXISTS {}".format(table))
5555

56+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
57+
"This is a MySQL test")
5658
def test_mysql_operator_test(self):
5759
sql = """
5860
CREATE TABLE IF NOT EXISTS test_airflow (
@@ -66,8 +68,11 @@ def test_mysql_operator_test(self):
6668
dag=self.dag)
6769
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
6870

71+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
72+
"This is a MySQL test")
6973
def test_mysql_operator_test_multi(self):
7074
sql = [
75+
"CREATE TABLE IF NOT EXISTS test_airflow (dummy VARCHAR(50))",
7176
"TRUNCATE TABLE test_airflow",
7277
"INSERT INTO test_airflow VALUES ('X')",
7378
]
@@ -79,6 +84,8 @@ def test_mysql_operator_test_multi(self):
7984
)
8085
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
8186

87+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
88+
"This is a MySQL test")
8289
def test_mysql_hook_test_bulk_load(self):
8390
records = ("foo", "bar", "baz")
8491

@@ -101,6 +108,8 @@ def test_mysql_hook_test_bulk_load(self):
101108
results = tuple(result[0] for result in c.fetchall())
102109
self.assertEqual(sorted(results), sorted(records))
103110

111+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
112+
"This is a MySQL test")
104113
def test_mysql_hook_test_bulk_dump(self):
105114
from airflow.hooks.mysql_hook import MySqlHook
106115
hook = MySqlHook('airflow_db')
@@ -112,6 +121,8 @@ def test_mysql_hook_test_bulk_dump(self):
112121
self.skipTest("Skip test_mysql_hook_test_bulk_load "
113122
"since file output is not permitted")
114123

124+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
125+
"This is a MySQL test")
115126
@mock.patch('airflow.hooks.mysql_hook.MySqlHook.get_conn')
116127
def test_mysql_hook_test_bulk_dump_mock(self, mock_get_conn):
117128
mock_execute = mock.MagicMock()
@@ -131,6 +142,8 @@ def test_mysql_hook_test_bulk_dump_mock(self, mock_get_conn):
131142
""".format(tmp_file=tmp_file, table=table)
132143
assertEqualIgnoreMultipleSpaces(self, mock_execute.call_args[0][0], query)
133144

145+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
146+
"This is a MySQL test")
134147
def test_mysql_to_mysql(self):
135148
sql = "SELECT * FROM INFORMATION_SCHEMA.TABLES LIMIT 100;"
136149
from airflow.operators.generic_transfer import GenericTransfer
@@ -148,6 +161,8 @@ def test_mysql_to_mysql(self):
148161
dag=self.dag)
149162
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
150163

164+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
165+
"This is a MySQL test")
151166
def test_overwrite_schema(self):
152167
"""
153168
Verifies option to overwrite connection schema
@@ -177,6 +192,16 @@ def setUp(self):
177192
dag = DAG(TEST_DAG_ID, default_args=args)
178193
self.dag = dag
179194

195+
def tearDown(self):
196+
tables_to_drop = ['test_postgres_to_postgres', 'test_airflow']
197+
from airflow.hooks.postgres_hook import PostgresHook
198+
with PostgresHook().get_conn() as conn:
199+
with conn.cursor() as cur:
200+
for t in tables_to_drop:
201+
cur.execute("DROP TABLE IF EXISTS {}".format(t))
202+
203+
@unittest.skipUnless('postgres' in configuration.conf.get('core', 'sql_alchemy_conn'),
204+
"This is a Postgres test")
180205
def test_postgres_operator_test(self):
181206
sql = """
182207
CREATE TABLE IF NOT EXISTS test_airflow (
@@ -197,8 +222,11 @@ def test_postgres_operator_test(self):
197222
end_date=DEFAULT_DATE,
198223
ignore_ti_state=True)
199224

225+
@unittest.skipUnless('postgres' in configuration.conf.get('core', 'sql_alchemy_conn'),
226+
"This is a Postgres test")
200227
def test_postgres_operator_test_multi(self):
201228
sql = [
229+
"CREATE TABLE IF NOT EXISTS test_airflow (dummy VARCHAR(50))",
202230
"TRUNCATE TABLE test_airflow",
203231
"INSERT INTO test_airflow VALUES ('X')",
204232
]
@@ -207,6 +235,8 @@ def test_postgres_operator_test_multi(self):
207235
task_id='postgres_operator_test_multi', sql=sql, dag=self.dag)
208236
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
209237

238+
@unittest.skipUnless('postgres' in configuration.conf.get('core', 'sql_alchemy_conn'),
239+
"This is a Postgres test")
210240
def test_postgres_to_postgres(self):
211241
sql = "SELECT * FROM INFORMATION_SCHEMA.TABLES LIMIT 100;"
212242
from airflow.operators.generic_transfer import GenericTransfer
@@ -224,6 +254,8 @@ def test_postgres_to_postgres(self):
224254
dag=self.dag)
225255
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
226256

257+
@unittest.skipUnless('postgres' in configuration.conf.get('core', 'sql_alchemy_conn'),
258+
"This is a Postgres test")
227259
def test_vacuum(self):
228260
"""
229261
Verifies the VACUUM operation runs well with the PostgresOperator
@@ -238,6 +270,8 @@ def test_vacuum(self):
238270
autocommit=True)
239271
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
240272

273+
@unittest.skipUnless('postgres' in configuration.conf.get('core', 'sql_alchemy_conn'),
274+
"This is a Postgres test")
241275
def test_overwrite_schema(self):
242276
"""
243277
Verifies option to overwrite connection schema
@@ -343,11 +377,15 @@ def tearDown(self):
343377
with MySqlHook().get_conn() as cur:
344378
cur.execute("DROP TABLE IF EXISTS baby_names CASCADE;")
345379

380+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
381+
"This is a MySQL test")
346382
def test_clear(self):
347383
self.dag.clear(
348384
start_date=DEFAULT_DATE,
349385
end_date=timezone.utcnow())
350386

387+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
388+
"This is a MySQL test")
351389
def test_mysql_to_hive(self):
352390
from airflow.operators.mysql_to_hive import MySqlToHiveTransfer
353391
sql = "SELECT * FROM baby_names LIMIT 1000;"
@@ -361,6 +399,8 @@ def test_mysql_to_hive(self):
361399
dag=self.dag)
362400
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
363401

402+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
403+
"This is a MySQL test")
364404
def test_mysql_to_hive_partition(self):
365405
from airflow.operators.mysql_to_hive import MySqlToHiveTransfer
366406
sql = "SELECT * FROM baby_names LIMIT 1000;"
@@ -376,6 +416,8 @@ def test_mysql_to_hive_partition(self):
376416
dag=self.dag)
377417
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
378418

419+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
420+
"This is a MySQL test")
379421
def test_mysql_to_hive_tblproperties(self):
380422
from airflow.operators.mysql_to_hive import MySqlToHiveTransfer
381423
sql = "SELECT * FROM baby_names LIMIT 1000;"
@@ -390,6 +432,8 @@ def test_mysql_to_hive_tblproperties(self):
390432
dag=self.dag)
391433
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
392434

435+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
436+
"This is a MySQL test")
393437
@mock.patch('airflow.hooks.hive_hooks.HiveCliHook.load_file')
394438
def test_mysql_to_hive_type_conversion(self, mock_load_file):
395439
mysql_table = 'test_mysql_to_hive'
@@ -433,6 +477,8 @@ def test_mysql_to_hive_type_conversion(self, mock_load_file):
433477
with m.get_conn() as c:
434478
c.execute("DROP TABLE IF EXISTS {}".format(mysql_table))
435479

480+
@unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
481+
"This is a MySQL test")
436482
def test_mysql_to_hive_verify_loaded_values(self):
437483
mysql_table = 'test_mysql_to_hive'
438484
hive_table = 'test_mysql_to_hive'
File renamed without changes.
File renamed without changes.

tests/jobs.py tests/test_jobs.py

File renamed without changes.

tests/utils.py tests/test_utils.py

File renamed without changes.

0 commit comments

Comments
 (0)