def test_hive_stats(self):
op = HiveStatsCollectionOperator(
task_id='hive_stats_check',
table="airflow.static_babynames_partitioned",
partition={'ds': DEFAULT_DATE_DS},
dag=self.dag)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
def test_runs_for_hive_stats(self, mock_hive_metastore_hook):
mock_mysql_hook = MockMySqlHook()
mock_presto_hook = MockPrestoHook()
with patch('airflow.operators.hive_stats_operator.PrestoHook',
return_value=mock_presto_hook):
with patch('airflow.operators.hive_stats_operator.MySqlHook',
return_value=mock_mysql_hook):
op = HiveStatsCollectionOperator(
task_id='hive_stats_check',
table="airflow.static_babynames_partitioned",
partition={'ds': DEFAULT_DATE_DS},
dag=self.dag)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
select_count_query = "SELECT COUNT(*) AS __count FROM airflow." \
+ "static_babynames_partitioned WHERE ds = '2015-01-01';"
mock_presto_hook.get_first.assert_called_with(hql=select_count_query)
expected_stats_select_query = "SELECT 1 FROM hive_stats WHERE table_name='airflow." \
+ "static_babynames_partitioned' AND " \
+ "partition_repr='{\"ds\": \"2015-01-01\"}' AND " \
+ "dttm='2015-01-01T00:00:00+00:00' " \
+ "LIMIT 1;"
raw_stats_select_query = mock_mysql_hook.get_records.call_args_list[0][0][0]
actual_stats_select_query = re.sub(r'\s{2,}', ' ', raw_stats_select_query).strip()
self.assertEqual(expected_stats_select_query, actual_stats_select_query)
insert_rows_val = [('2015-01-01', '2015-01-01T00:00:00+00:00',
'airflow.static_babynames_partitioned',
'{"ds": "2015-01-01"}', '', 'count', ['val_0', 'val_1'])]
mock_mysql_hook.insert_rows.assert_called_with(table='hive_stats',
rows=insert_rows_val,
target_fields=[
'ds',
'dttm',
'table_name',
'partition_repr',
'col',
'metric',
'value',
])