def test_push_from_local_should_delete_file_if_exists_and_save_file(
self, base_conn_mock, samba_hook_mock):
base_conn_mock.return_value = connection
samba_hook_mock.get_conn.return_value = mock.Mock()
samba_hook_mock.return_value.exists.return_value = False
samba_hook_mock.return_value.exists.return_value = False
hook = SambaHook('samba_default')
destination_folder = "/path/to/dest"
destination_filepath = destination_folder + "/file"
local_filepath = "/path/to/local/file"
hook.push_from_local(destination_filepath=destination_filepath,
local_filepath=local_filepath)
base_conn_mock.assert_called_once_with('samba_default')
samba_hook_mock.assert_called_once()
samba_hook_mock.return_value.exists.assert_has_calls(
[call(destination_filepath),
call(destination_folder)])
samba_hook_mock.return_value.isfile.assert_not_called()
samba_hook_mock.return_value.remove.assert_not_called()
samba_hook_mock.return_value.mkdir.assert_called_once_with(
destination_folder)
samba_hook_mock.return_value.upload.assert_called_once_with(
local_filepath, destination_filepath)
def execute(self, context):
with NamedTemporaryFile() as tmp_file:
self.log.info("Fetching file from Hive")
hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id)
hive.to_csv(hql=self.hql, csv_filepath=tmp_file.name, hive_conf=context_to_airflow_vars(context))
self.log.info("Pushing to samba")
samba = SambaHook(samba_conn_id=self.samba_conn_id)
samba.push_from_local(self.destination_filepath, tmp_file.name)
def test_push_from_local_should_succeed_if_destination_has_same_name_but_not_a_file(
self, base_conn_mock, samba_hook_mock
):
base_conn_mock.return_value = connection
samba_hook_mock.get_conn.return_value = mock.Mock()
samba_hook_mock.return_value.exists.return_value = True
samba_hook_mock.return_value.isfile.return_value = False
samba_hook_mock.return_value.exists.return_value = True
hook = SambaHook('samba_default')
destination_filepath = "/path/to/dest/file"
local_filepath = "/path/to/local/file"
hook.push_from_local(destination_filepath=destination_filepath, local_filepath=local_filepath)
base_conn_mock.assert_called_once_with('samba_default')
samba_hook_mock.assert_called_once()
samba_hook_mock.return_value.exists.assert_called_once_with(destination_filepath)
samba_hook_mock.return_value.isfile.assert_called_once_with(destination_filepath)
samba_hook_mock.return_value.remove.assert_not_called()
samba_hook_mock.return_value.upload.assert_called_once_with(local_filepath, destination_filepath)
def test_get_conn(self, get_conn_mock):
get_conn_mock.return_value = connection
hook = SambaHook('samba_default')
self.assertEqual(smbclient.SambaClient, type(hook.get_conn()))
get_conn_mock.assert_called_once_with('samba_default')
def test_get_conn_should_fail_if_conn_id_does_not_exist(self):
with self.assertRaises(AirflowException):
SambaHook('conn')
def test_get_conn(self, get_conn_mock):
get_conn_mock.return_value = connection
hook = SambaHook('samba_default')
assert isinstance(hook.get_conn(), smbclient.SambaClient)
get_conn_mock.assert_called_once_with('samba_default')
