def test_key_and_connection(self): from azure.storage.file import FileService hook = AzureFileShareHook(wasb_conn_id='wasb_test_key') self.assertEqual(hook.conn_id, 'wasb_test_key') self.assertIsNone(hook._conn) self.assertIsInstance(hook.get_conn(), FileService)
def test_get_file_to_stream(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.get_file_to_stream('stream', 'share', 'directory', 'file', max_connections=1) mock_instance.get_file_to_stream.assert_called_once_with( 'share', 'directory', 'file', 'stream', max_connections=1 )
def test_load_string(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.load_string('big string', 'share', 'directory', 'file', timeout=1) mock_instance.create_file_from_text.assert_called_once_with( 'share', 'directory', 'file', 'big string', timeout=1 )
def test_load_file(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.load_file('path', 'share', 'directory', 'file', max_connections=1) mock_instance.create_file_from_path.assert_called_once_with( 'share', 'directory', 'file', 'path', max_connections=1 )
def test_create_directory(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.create_directory('share', 'directory', timeout=1) mock_instance.create_directory.assert_called_once_with('share', 'directory', timeout=1)
def test_key_and_connection(self): from azure.storage.file import FileService hook = AzureFileShareHook(wasb_conn_id='wasb_test_key') assert hook.conn_id == 'wasb_test_key' assert hook._conn is None assert isinstance(hook.get_conn(), FileService)
def test_list_directories_and_files(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.list_directories_and_files('share', 'directory', timeout=1) mock_instance.list_directories_and_files.assert_called_once_with( 'share', 'directory', timeout=1 )
def test_check_for_directory(self, mock_service): mock_instance = mock_service.return_value mock_instance.exists.return_value = True hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') assert hook.check_for_directory('share', 'directory', timeout=3) mock_instance.exists.assert_called_once_with('share', 'directory', timeout=3)
def test_list_files(self, mock_service): mock_instance = mock_service.return_value mock_instance.list_directories_and_files.return_value = [ File("file1"), File("file2"), Directory("dir1"), Directory("dir2"), ] hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') files = hook.list_files('share', 'directory', timeout=1) self.assertEqual(files, ["file1", 'file2']) mock_instance.list_directories_and_files.assert_called_once_with('share', 'directory', timeout=1)
def upload_file_from_string(cls, string_data: str, share_name: str, wasb_conn_id: str, file_name: str, directory: Optional[str] = None): hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id) hook.load_string( string_data=string_data, share_name=share_name, directory_name=directory, file_name=file_name, )
def create_fileshare(): """Create a fileshare with directory""" hook = AzureFileShareHook() hook.create_share(NAME) hook.create_directory(share_name=NAME, directory_name=DIRECTORY) exists = hook.check_for_directory(share_name=NAME, directory_name=DIRECTORY) if not exists: raise Exception
def test_sas_token(self): from azure.storage.file import FileService hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') self.assertEqual(hook.conn_id, 'wasb_test_sas_token') self.assertIsInstance(hook.get_conn(), FileService)
def test_sas_token(self): from azure.storage.file import FileService hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') assert hook.conn_id == 'wasb_test_sas_token' assert isinstance(hook.get_conn(), FileService)
def execute(self, context): azure_fileshare_hook = AzureFileShareHook(self.wasb_conn_id) files = azure_fileshare_hook.list_files( share_name=self.share_name, directory_name=self.directory_name) gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, impersonation_chain=self.google_impersonation_chain, ) dest_gcs_bucket, dest_gcs_object_prefix = _parse_gcs_url(self.dest_gcs) # pylint: disable=too-many-nested-blocks if not self.replace: # if we are not replacing -> list all files in the GCS bucket # and only keep those files which are present in # S3 and not in Google Cloud Storage existing_files_prefixed = gcs_hook.list( dest_gcs_bucket, prefix=dest_gcs_object_prefix) existing_files = [] # Remove the object prefix itself, an empty directory was found if dest_gcs_object_prefix in existing_files_prefixed: existing_files_prefixed.remove(dest_gcs_object_prefix) # Remove the object prefix from all object string paths for file in existing_files_prefixed: if file.startswith(dest_gcs_object_prefix): existing_files.append(file[len(dest_gcs_object_prefix):]) else: existing_files.append(file) files = list(set(files) - set(existing_files)) if files: self.log.info('%s files are going to be synced.', len(files)) else: self.log.info('There are no new files to sync. Have a nice day!') for file in files: with NamedTemporaryFile() as temp_file: azure_fileshare_hook.get_file_to_stream( stream=temp_file, share_name=self.share_name, directory_name=self.directory_name, file_name=file, ) temp_file.flush() # There will always be a '/' before file because it is # enforced at instantiation time dest_gcs_object = dest_gcs_object_prefix + file gcs_hook.upload(dest_gcs_bucket, dest_gcs_object, temp_file.name, gzip=self.gzip) if files: self.log.info( "All done, uploaded %d files to Google Cloud Storage.", len(files)) else: self.log.info( 'In sync, no files needed to be uploaded to Google Cloud Storage' ) return files
def delete_share(cls, share_name: str, wasb_conn_id: str): hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id) hook.delete_share(share_name)
def create_directory(cls, share_name: str, wasb_conn_id: str, directory: str): hook = AzureFileShareHook(wasb_conn_id=wasb_conn_id) hook.create_directory(share_name=share_name, directory_name=directory)
def test_delete_share(self, mock_service): mock_instance = mock_service.return_value hook = AzureFileShareHook(wasb_conn_id='wasb_test_sas_token') hook.delete_share('my_share') mock_instance.delete_share.assert_called_once_with('my_share')
def delete_fileshare(): """Delete a fileshare""" hook = AzureFileShareHook() hook.delete_share(NAME)