def test_hive2samba(self):
op = Hive2SambaOperator(
task_id='hive2samba_check',
samba_conn_id='tableau_samba',
hql="SELECT * FROM airflow.static_babynames LIMIT 10000",
destination_filepath='test_airflow.csv',
dag=self.dag)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
def test_execute(self, mock_tmp_file, mock_hive_hook, mock_samba_hook):
type(mock_tmp_file).name = PropertyMock(return_value='tmp_file')
mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file)
context = {}
Hive2SambaOperator(**self.kwargs).execute(context)
mock_hive_hook.assert_called_once_with(
hiveserver2_conn_id=self.kwargs['hiveserver2_conn_id'])
mock_hive_hook.return_value.to_csv.assert_called_once_with(
hql=self.kwargs['hql'],
csv_filepath=mock_tmp_file.name,
hive_conf=context_to_airflow_vars(context))
mock_samba_hook.assert_called_once_with(
samba_conn_id=self.kwargs['samba_conn_id'])
mock_samba_hook.return_value.push_from_local.assert_called_once_with(
self.kwargs['destination_filepath'], mock_tmp_file.name)
def test_hive2samba(self, mock_hive_server_hook, mock_temp_dir):
mock_temp_dir.return_value = "tst"
samba_hook = MockSambaHook(self.kwargs['samba_conn_id'])
samba_hook.upload = MagicMock()
with patch('airflow.operators.hive_to_samba_operator.SambaHook',
return_value=samba_hook):
samba_hook.conn.upload = MagicMock()
op = Hive2SambaOperator(
task_id='hive2samba_check',
samba_conn_id='tableau_samba',
hql="SELECT * FROM airflow.static_babynames LIMIT 10000",
destination_filepath='test_airflow.csv',
dag=self.dag)
op.run(start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE,
ignore_ti_state=True)
samba_hook.conn.upload.assert_called_with('/tmp/tmptst',
'test_airflow.csv')