def test_hive_to_mysql(self):
test_hive_results = 'test_hive_results'
mock_hive_hook = MockHiveServer2Hook()
mock_hive_hook.get_records = MagicMock(return_value=test_hive_results)
mock_mysql_hook = MockMySqlHook()
mock_mysql_hook.run = MagicMock()
mock_mysql_hook.insert_rows = MagicMock()
with patch('airflow.operators.hive_to_mysql.HiveServer2Hook',
return_value=mock_hive_hook):
with patch('airflow.operators.hive_to_mysql.MySqlHook',
return_value=mock_mysql_hook):
op = HiveToMySqlTransfer(
mysql_conn_id='airflow_db',
task_id='hive_to_mysql_check',
sql="""
SELECT name
FROM airflow.static_babynames
LIMIT 100
""",
mysql_table='test_static_babynames',
mysql_preoperator=[
'DROP TABLE IF EXISTS test_static_babynames;',
'CREATE TABLE test_static_babynames (name VARCHAR(500))',
],
dag=self.dag)
op.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
op.run(start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE,
ignore_ti_state=True)
raw_select_name_query = mock_hive_hook.get_records.call_args_list[0][
0][0]
actual_select_name_query = re.sub(r'\s{2,}', ' ',
raw_select_name_query).strip()
expected_select_name_query = 'SELECT name FROM airflow.static_babynames LIMIT 100'
self.assertEqual(expected_select_name_query, actual_select_name_query)
actual_hive_conf = mock_hive_hook.get_records.call_args_list[0][1][
'hive_conf']
expected_hive_conf = {
'airflow.ctx.dag_owner': 'airflow',
'airflow.ctx.dag_id': 'test_dag_id',
'airflow.ctx.task_id': 'hive_to_mysql_check',
'airflow.ctx.execution_date': '2015-01-01T00:00:00+00:00'
}
self.assertEqual(expected_hive_conf, actual_hive_conf)
expected_mysql_preoperator = [
'DROP TABLE IF EXISTS test_static_babynames;',
'CREATE TABLE test_static_babynames (name VARCHAR(500))'
]
mock_mysql_hook.run.assert_called_with(expected_mysql_preoperator)
mock_mysql_hook.insert_rows.assert_called_with(
table='test_static_babynames', rows=test_hive_results)
def test_hive_to_mysql(self):
op = HiveToMySqlTransfer(
mysql_conn_id='airflow_db',
task_id='hive_to_mysql_check',
create=True,
sql="""
SELECT name
FROM airflow.static_babynames
LIMIT 100
""",
mysql_table='test_static_babynames',
mysql_preoperator=[
'DROP TABLE IF EXISTS test_static_babynames;',
'CREATE TABLE test_static_babynames (name VARCHAR(500))',
],
dag=self.dag)
op.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)