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')