from ftplib import error_perm import os from unittest.mock import MagicMock, call from unittest.mock import mock_open from unittest.mock import patch from garcon_contrib.ftp import garcon_ftp # noqa import paramiko import pytest US_SFTP_CREDS = dict( username='user', host='host', port=22) CANADA_FTP_CREDS = dict( username='user', password='password', host='host', port=21) US_FILE_NAMES_LIST = [ 'TRACK_ORCH.txt.gz', 'PRODUCT_TRACK_ORCH.txt.gz', 'PRODUCT_ORCH.txt.gz'] CANADA_FILE_NAMES_LIST = [ 'ORCH_Canada_Track.txt.gz', 'ORCH_Canada_Product.txt.gz'] US_PRIVATE_KEY_PATH = './id_rsa_soundscan_us' def test_upload_from_local_to_ftp(monkeypatch): # Mock sftp client object sftp_obj = MagicMock() monkeypatch.setattr(sftp_obj, 'chdir', MagicMock(return_value='')) monkeypatch.setattr(sftp_obj, 'mkdir', MagicMock(return_value='')) monkeypatch.setattr(sftp_obj, 'put', MagicMock(return_value='')) monkeypatch.setattr(sftp_obj, 'close', MagicMock(return_value='')) monkeypatch.setattr( paramiko.SFTPClient, 'from_transport', MagicMock(return_value=sftp_obj)) # Mock transport object transport_obj = MagicMock() monkeypatch.setattr(transport_obj, 'connect', MagicMock(return_value='')) monkeypatch.setattr(transport_obj, 'close', MagicMock(return_value='')) monkeypatch.setattr( paramiko, 'Transport', MagicMock(return_value=transport_obj)) # Mock paramiko.RSAKey object monkeypatch.setattr( paramiko.RSAKey, 'from_private_key_file', MagicMock(return_value='')) garcon_ftp._upload_from_local_to_ftp( US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir', 'test_key_path') assert not sftp_obj.mkdir.called assert sftp_obj.chdir.called assert sftp_obj.put.called assert sftp_obj.close.called # Test if IOError exception thrown monkeypatch.setattr( sftp_obj, 'chdir', MagicMock(return_value='', side_effect=IOError('Boom!'))) with pytest.raises(IOError): garcon_ftp._upload_from_local_to_ftp( US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir', 'test_key_path') assert sftp_obj.mkdir.called # Mock regular ftp client object ftp_obj = MagicMock() monkeypatch.setattr(garcon_ftp, 'FTP', MagicMock(return_value=ftp_obj)) m = mock_open() with patch('builtins.open', m, create=True): garcon_ftp._upload_from_local_to_ftp( CANADA_FTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir') assert ftp_obj.mkd.called assert ftp_obj.close.called assert ftp_obj.storbinary.called @patch('os.stat') def test_download_from_sftp(mock_stat, monkeypatch): mock_stat().st_size = 5 copyfile_obj = MagicMock() monkeypatch.setattr(garcon_ftp.shutil, 'copyfileobj', copyfile_obj) # Mock sftp client object sftp_obj = MagicMock() sftp_obj.open = mock_open() sftp_obj.stat.return_value.st_size = 5 monkeypatch.setattr( paramiko.SFTPClient, 'from_transport', MagicMock(return_value=sftp_obj)) # Mock transport object transport_obj = MagicMock() monkeypatch.setattr( paramiko, 'Transport', MagicMock(return_value=transport_obj)) # Mock paramiko.RSAKey object monkeypatch.setattr( paramiko.RSAKey, 'from_private_key_file', MagicMock(return_value='test_rsa_key')) open_local_obj = mock_open() with patch( 'garcon_contrib.ftp.garcon_ftp.open', open_local_obj, create=True): # assert sftp call with a RSA key & password garcon_ftp._download_from_sftp( US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir', 'test_key_path') transport_obj.connect.assert_called_with( username='user', password=None, pkey='test_rsa_key') sftp_obj.chdir.assert_called_with('test_remote_dir') open_local_obj.assert_called_with('test_local_dir/test', 'wb') sftp_obj.open.assert_called_with('test') copyfile_obj.assert_called_with(sftp_obj.open(), open_local_obj()) sftp_obj.close.assert_called_with() # assert sftp call with just a password (no RSA key) transport_obj.reset_mock() sftp_creds = dict( username='user', password='password', host='host', port=22) garcon_ftp._download_from_sftp( sftp_creds, 'test', 'test_remote_dir', 'test_local_dir') transport_obj.connect.assert_called_with( username='user', password='password', pkey=None) # assert IOError raised if sftp file size mismatch occurs sftp_obj.stat.return_value.st_size = 7 with pytest.raises(IOError): garcon_ftp._download_from_sftp( sftp_creds, 'test', 'test_remote_dir', 'test_local_dir') def test_download_from_ftp(monkeypatch): # Mock regular ftp client object ftp_obj = MagicMock() monkeypatch.setattr(garcon_ftp, 'FTP', MagicMock(return_value=ftp_obj)) m = mock_open() with patch('builtins.open', m, create=True): garcon_ftp._download_from_ftp( CANADA_FTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir') assert ftp_obj.close.called assert ftp_obj.retrbinary.called @patch.object(garcon_ftp.boto3, "resource") @patch.object(garcon_ftp, "_upload_from_local_to_ftp") @patch("builtins.open") def test_copy_from_s3_to_ftp( open_mock, upload_ftp_mock, resource_mock, monkeypatch): monkeypatch.setattr( garcon_ftp, '_upload_from_local_to_ftp', MagicMock(return_value='')) monkeypatch.setattr( os, 'remove', MagicMock(return_value='')) garcon_ftp.copy_from_s3_to_ftp( MagicMock(), US_SFTP_CREDS, US_FILE_NAMES_LIST, 's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01', None) bucket_mock = resource_mock.return_value.Bucket.return_value assert open_mock.call_args_list == [ call('./TRACK_ORCH.txt.gz', 'wb'), call('./PRODUCT_TRACK_ORCH.txt.gz', 'wb'), call('./PRODUCT_ORCH.txt.gz', 'wb'), ] open_mock_return = open_mock.return_value.__enter__.return_value assert bucket_mock.download_fileobj.call_args_list == [ call('bucket_path/TRACK_ORCH.txt.gz', open_mock_return), call('bucket_path/PRODUCT_TRACK_ORCH.txt.gz', open_mock_return), call('bucket_path/PRODUCT_ORCH.txt.gz', open_mock_return), ] assert garcon_ftp._upload_from_local_to_ftp.call_args_list == [ call(US_SFTP_CREDS, 'TRACK_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', '.', None), call(US_SFTP_CREDS, 'PRODUCT_TRACK_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', '.', None), call(US_SFTP_CREDS, 'PRODUCT_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', '.', None), ] assert os.remove.call_args_list == [ call('./TRACK_ORCH.txt.gz'), call('./PRODUCT_TRACK_ORCH.txt.gz'), call('./PRODUCT_ORCH.txt.gz') ] @patch.object(garcon_ftp.boto3, "resource") def test_copy_from_ftp_to_s3(resource_mock, monkeypatch): monkeypatch.setattr(garcon_ftp, '_download_from_sftp', MagicMock()) monkeypatch.setattr(garcon_ftp, '_download_from_ftp', MagicMock()) monkeypatch.setattr(os, 'remove', MagicMock(return_value='')) m = mock_open() with patch('builtins.open', m, create=True): # Test sftp calls made when sftp creds are provided response = garcon_ftp.copy_from_ftp_to_s3( MagicMock(), US_SFTP_CREDS, US_FILE_NAMES_LIST, 's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01', 'key') garcon_ftp._download_from_sftp.assert_any_call( US_SFTP_CREDS, 'TRACK_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', './', 'key') garcon_ftp._download_from_sftp.assert_any_call( US_SFTP_CREDS, 'PRODUCT_TRACK_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', './', 'key') garcon_ftp._download_from_sftp.assert_any_call( US_SFTP_CREDS, 'PRODUCT_ORCH.txt.gz', '/orchard/soundscan_us/2015-10-01/', './', 'key') assert len(response['files']) == 3 for resp in response['files']: assert resp['status'] is True assert resp.get('exception') is None assert os.remove.called os.remove.reset_mock() # Test ftp calls made when ftp creds are provided response = garcon_ftp.copy_from_ftp_to_s3( MagicMock(), CANADA_FTP_CREDS, CANADA_FILE_NAMES_LIST, 's3://bucket_name/bucket_path', '/orchard/soundscan_canada/2015-10-01', None) garcon_ftp._download_from_ftp.assert_any_call( CANADA_FTP_CREDS, 'ORCH_Canada_Track.txt.gz', '/orchard/soundscan_canada/2015-10-01/', './') garcon_ftp._download_from_ftp.assert_any_call( CANADA_FTP_CREDS, 'ORCH_Canada_Product.txt.gz', '/orchard/soundscan_canada/2015-10-01/', './') assert len(response['files']) == 2 for resp in response['files']: assert resp['status'] is True assert resp.get('exception') is None assert os.remove.called def _download_from_ftp( sftp_creds, file_name, remote_ftp_dir, local_dir): raise error_perm('No download for you') def _download_from_sftp( sftp_creds, file_name, remote_ftp_dir, local_dir, pkey): raise FileNotFoundError('No download for you') @patch.object(garcon_ftp.boto3, "resource") def test_copy_from_ftp_to_s3_fails(resource_mock, monkeypatch): monkeypatch.setattr( garcon_ftp, '_download_from_ftp', _download_from_ftp) monkeypatch.setattr( garcon_ftp, '_download_from_sftp', _download_from_sftp) monkeypatch.setattr( os, 'remove', MagicMock(return_value='')) response = garcon_ftp.copy_from_ftp_to_s3( MagicMock(), US_SFTP_CREDS, US_FILE_NAMES_LIST, 's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01', None) assert len(response['files']) == 3 for resp in response['files']: assert resp['status'] is False assert resp['exception'] is not None assert isinstance(resp['exception'], FileNotFoundError) assert not os.remove.called assert not resource_mock.upload_fileobj.called response = garcon_ftp.copy_from_ftp_to_s3( MagicMock(), CANADA_FTP_CREDS, CANADA_FILE_NAMES_LIST, 's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01', None) assert len(response['files']) == 2 for resp in response['files']: assert resp['status'] is False assert resp['exception'] is not None assert isinstance(resp['exception'], error_perm) assert not os.remove.called assert not resource_mock.upload_fileobj.called @patch('garcon_contrib.ftp.garcon_ftp.paramiko') def test_file_exist_on_sftp(mock_paramiko): mock_sftp_client = MagicMock() mock_sftp_client.stat.return_value = None mock_sftp = MagicMock() mock_sftp.from_transport.return_value = mock_sftp_client mock_transport = MagicMock() mock_transport.connect.return_value = 'connect is called' mock_paramiko.SFTPClient.from_transport.return_value = mock_sftp_client mock_paramiko.Transport.return_value = mock_transport sftp_creds = dict( host='host', port=22, username='user', password='password' ) actual = garcon_ftp.file_exist_on_sftp('test_path', **sftp_creds) assert mock_sftp_client.stat.called assert actual @patch('garcon_contrib.ftp.garcon_ftp.paramiko') def test_file_exist_on_sftp_file_not_found(mock_paramiko): mock_sftp_client = MagicMock() mock_sftp_client.stat.return_value = None mock_sftp_client.stat.side_effect = FileNotFoundError('No such file') mock_sftp = MagicMock() mock_sftp.from_transport.return_value = mock_sftp_client mock_transport = MagicMock() mock_transport.connect.return_value = 'connect is called' mock_paramiko.SFTPClient.from_transport.return_value = mock_sftp_client mock_paramiko.Transport.return_value = mock_transport sftp_creds = dict( host='host', port=22, username='user', password='password' ) actual = garcon_ftp.file_exist_on_sftp('test_path', **sftp_creds) assert mock_sftp_client.stat.called assert not actual @patch.object(garcon_ftp.boto3, "resource") def test_copy_file_from_ftp_to_s3(resource_mock, monkeypatch): monkeypatch.setattr(garcon_ftp, '_download_from_sftp', MagicMock()) monkeypatch.setattr(garcon_ftp, '_download_from_ftp', MagicMock()) monkeypatch.setattr(os, 'remove', MagicMock(return_value='')) activity = MagicMock() # Test sftp calls made when sftp creds are provided ftp_path = '/ftp/path' ftp_file_name = 'ftp_file' s3_path = 's3://bucket_name/bucket_path' s3_file_name = 's3_file' pkey = 'key' with patch('builtins.open') as open_mock: sftp_response = garcon_ftp.copy_file_from_ftp_to_s3( activity, US_SFTP_CREDS, ftp_path, ftp_file_name, s3_path, s3_file_name, pkey) assert garcon_ftp._download_from_sftp.call_args_list == [ call(US_SFTP_CREDS, ftp_file_name, ftp_path, '.', pkey) ] assert open_mock.call_args_list == [ call('./ftp_file', 'rb') ] assert sftp_response == { 'file': 's3_file', 'file_size': (resource_mock.return_value. Object.return_value.content_length), 'status': True } assert os.remove.call_args_list == [ call('./ftp_file') ] # TODO: extract below as paramaterized test case # instead of reset mock os.remove.reset_mock() # Test ftp calls made when ftp creds are provided ftp_response = garcon_ftp.copy_file_from_ftp_to_s3( activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path, s3_file_name, None) assert garcon_ftp._download_from_ftp.call_args_list == [ call(CANADA_FTP_CREDS, ftp_file_name, ftp_path, '.') ] assert ftp_response == { 'file': 's3_file', 'file_size': (resource_mock.return_value. Object.return_value.content_length), 'status': True } assert os.remove.call_args_list == [ call('./ftp_file') ] # Test target file name file_response = garcon_ftp.copy_file_from_ftp_to_s3( activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path) assert file_response['file'] == ftp_file_name @patch.object(garcon_ftp.boto3, "resource") def test_copy_file_from_ftp_to_s3_fails(resource_mock, monkeypatch): monkeypatch.setattr(garcon_ftp, '_download_from_sftp', _download_from_sftp) monkeypatch.setattr(garcon_ftp, '_download_from_ftp', _download_from_ftp) monkeypatch.setattr(os, 'remove', MagicMock(return_value='')) activity = MagicMock() # Test sftp calls made when sftp creds are provided ftp_path = '/ftp/path' ftp_file_name = 'ftp_file' s3_path = 's3://bucket_name/bucket_path' s3_file_name = 's3_file' pkey = 'key' sftp_response = garcon_ftp.copy_file_from_ftp_to_s3( activity, US_SFTP_CREDS, ftp_path, ftp_file_name, s3_path, s3_file_name, pkey) assert sftp_response['status'] is False assert sftp_response['exception'] is not None assert isinstance(sftp_response['exception'], FileNotFoundError) assert not os.remove.called assert not resource_mock.upload_fileobj.called # Test ftp calls made when ftp creds are provided ftp_response = garcon_ftp.copy_file_from_ftp_to_s3( activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path, s3_file_name, None) assert ftp_response['status'] is False assert ftp_response['exception'] is not None assert isinstance(ftp_response['exception'], error_perm) assert not os.remove.called assert not resource_mock.upload_fileobj.called @patch('garcon_contrib.ftp.garcon_ftp.paramiko') def test_get_list_of_files_and_directories_sftp(mock_paramiko): """Test _get_list_of_files_and_directories_sftp method.""" mock_sftp_client = MagicMock() mock_sftp_client.listdir_iter.return_value = [ MagicMock(filename='file1'), MagicMock(filename='file2'), MagicMock(filename='test')] mock_sftp = MagicMock() mock_sftp.from_transport.return_value = mock_sftp_client mock_transport = MagicMock() mock_transport.connect.return_value = 'connect is called' mock_paramiko.SFTPClient.from_transport.return_value.__enter__. \ return_value = mock_sftp_client mock_paramiko.Transport.return_value.__enter__.return_value = \ mock_transport sftp_creds = dict( host='host', port=22, username='user', password='password' ) actual = garcon_ftp._get_list_of_files_and_directories_sftp( sftp_creds, 'test_path', 'file') assert mock_sftp_client.listdir_iter.called assert actual == ['file1', 'file2'] # call without file_pattern actual = garcon_ftp._get_list_of_files_and_directories_sftp( sftp_creds, 'test_path') assert actual == ['file1', 'file2', 'test'] @patch('garcon_contrib.ftp.garcon_ftp.paramiko') def test_test_get_list_of_files_and_directories_sftp_path_not_found( mock_paramiko): """Test _get_list_of_files_and_directories_sftp method.""" mock_sftp_client = MagicMock() mock_sftp_client.listdir_iter.return_value = None mock_sftp_client.listdir_iter.side_effect = FileNotFoundError( 'No such file') mock_sftp = MagicMock() mock_sftp.from_transport.return_value = mock_sftp_client mock_transport = MagicMock() mock_transport.connect.return_value = 'connect is called' mock_paramiko.SFTPClient.from_transport.return_value.__enter__. \ return_value = mock_sftp_client mock_paramiko.Transport.return_value.__enter__.return_value = \ mock_transport sftp_creds = dict( host='host', port=22, username='user', password='password' ) with pytest.raises(FileNotFoundError): actual = garcon_ftp._get_list_of_files_and_directories_sftp( sftp_creds, 'test_path') assert mock_sftp_client.listdir_iter.called assert not actual def test_get_list_of_files_and_directories_ftp(monkeypatch): """Test _get_list_of_files_and_directories_ftp method.""" # Mock regular ftp client object ftp_obj = MagicMock() monkeypatch.setattr(garcon_ftp, 'FTP', MagicMock(return_value=ftp_obj)) ftp_obj.nlst.return_value = ['/dir/file1', '/dir/file2', 'dir/test'] ftp_creds = dict( host='host', port=21, username='user', password='password' ) actual = garcon_ftp._get_list_of_files_and_directories_ftp( ftp_creds, 'test_path', 'file') assert actual == ['file1', 'file2'] ftp_obj.nlst.assert_called_once_with('test_path') # call without file_pattern actual = garcon_ftp._get_list_of_files_and_directories_ftp( ftp_creds, 'test_path') assert actual == ['file1', 'file2', 'test'] @patch('garcon_contrib.ftp.garcon_ftp._get_list_of_files_and_directories_sftp') @patch('garcon_contrib.ftp.garcon_ftp._get_list_of_files_and_directories_ftp') def test_get_list_of_files_and_directories_call_with_sftp(mock_ftp, mock_sftp): """Test get_list_of_files_and_directories method with sftp creds.""" sftp_creds = dict( host='host', port=22, username='user', password='password' ) garcon_ftp.get_list_of_files_and_directories(sftp_creds, '') mock_ftp.assert_not_called() mock_sftp.assert_called_with(sftp_creds, '', None) sftp_creds = dict( host='host', port=22, username='user', pkey='/path/key.txt' ) garcon_ftp.get_list_of_files_and_directories(sftp_creds, '') mock_ftp.assert_not_called() mock_sftp.assert_called_with(sftp_creds, '', None) @patch('garcon_contrib.ftp.garcon_ftp._get_list_of_files_and_directories_sftp') @patch('garcon_contrib.ftp.garcon_ftp._get_list_of_files_and_directories_ftp') def test_get_list_of_files_and_directories_call_with_ftp(mock_ftp, mock_sftp): """Test get_list_of_files_and_directories method with ftp creds.""" ftp_creds = dict( host='host', port=21, username='user', password='password' ) garcon_ftp.get_list_of_files_and_directories(ftp_creds, '') mock_sftp.assert_not_called() mock_ftp.assert_called_with(ftp_creds, '', None)