"""Unit tests for the FTP or SFTP related tasks.""" from unittest.mock import MagicMock from unittest.mock import patch import pytest from feed_ingestion.tasks import ftp_amazon_marketshare_tasks @pytest.fixture() def sftp_conf_mock(): """Mock sftp configuration.""" return { 'account_name1': { 'host': 'host1', 'username': 'username1', 'password': 'password1', 'port': 22, }, 'account_name2': { 'host': 'host2', 'username': 'username2', 'password': 'password2', 'port': 22, }, } @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.FTP') def test_get_files_from_ftp(mock_ftp, sftp_conf_mock): """Test _get_files_from_ftp method.""" def sf_nlst(path): """Return list of files and directories.""" if path == '/root_path/': return ['AT', 'US'] return [ '{}/{}'.format(path, file) for file in [ 'pattern1.txt', 'pattern2.txt']] ftp_conn = MagicMock() ftp_conn.nlst.side_effect = sf_nlst mock_ftp.return_value = ftp_conn result = ftp_amazon_marketshare_tasks._get_files_from_ftp( sftp_conf_mock['account_name1'], '/root_path/{country_code}/rest_of_path', 'pattern') expected_response = { '/root_path/AT/rest_of_path/pattern1.txt': 'AT pattern1.txt', '/root_path/AT/rest_of_path/pattern2.txt': 'AT pattern2.txt', '/root_path/US/rest_of_path/pattern1.txt': 'US pattern1.txt', '/root_path/US/rest_of_path/pattern2.txt': 'US pattern2.txt'} assert result == expected_response @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.paramiko') def test_get_files_from_sftp(mock_paramiko, sftp_conf_mock): """Test _get_files_from_sftp method.""" def sf_listdir_attr(path): """Return list of files and directories.""" if path == '/root_path/': return ['AT', 'US'] return [ '{}/ {}'.format(path, file) for file in [ 'pattern1.txt', 'pattern2.txt']] mock_sftp_client = MagicMock() mock_sftp_client.listdir_attr.return_value = None mock_sftp_client.listdir_attr.side_effect = sf_listdir_attr 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 result = ftp_amazon_marketshare_tasks._get_files_from_sftp( sftp_conf_mock['account_name1'], '/root_path/{country_code}/rest_of_path', 'pattern') expected_response = { '/root_path/AT/rest_of_path/pattern1.txt': 'AT pattern1.txt', '/root_path/AT/rest_of_path/pattern2.txt': 'AT pattern2.txt', '/root_path/US/rest_of_path/pattern1.txt': 'US pattern1.txt', '/root_path/US/rest_of_path/pattern2.txt': 'US pattern2.txt'} assert result == expected_response @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks' '._get_files_from_sftp') @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.task_status') @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.' 'garcon_feed_status.delete_status') def test_check_for_new_files_if_there_are_new_files( mock_delete_status, mock_task_status, mock_get_files_from_sftp, sftp_conf_mock): """Test check_for_new_files if there are new files.""" response = { '/root_path/AT/rest_of_path/pattern1.txt': 'AT pattern1.txt', '/root_path/AT/rest_of_path/pattern2.txt': 'AT pattern2.txt', '/root_path/US/rest_of_path/pattern1.txt': 'US pattern1.txt', '/root_path/US/rest_of_path/pattern2.txt': 'US pattern2.txt'} mock_get_files_from_sftp.return_value = response mock_task_status.get_values.return_value = [] result = ftp_amazon_marketshare_tasks.check_for_new_files( activity=MagicMock(), feed_name='feed_name', date='2016-01-01', ftp_creds=sftp_conf_mock, ftp_path_template='/root_path/{country_code}/rest_of_path', file_pattern='pattern') expected_response = { 'account_name1': response, 'account_name2': response, } assert result == {'files_on_ftp': expected_response} mock_delete_status.assert_called_once_with('feed_name', '2016-01-01') @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks' '._get_files_from_sftp') @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.task_status') @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.' 'garcon_feed_status.delete_status') def test_check_for_new_files_if_there_are_no_new_files( mock_delete_status, mock_task_status, mock_get_files_from_sftp, sftp_conf_mock): """Test check_for_new_files if there are no new files.""" response = { '/root_path/AT/rest_of_path/pattern1.txt': 'AT pattern1.txt', '/root_path/AT/rest_of_path/pattern2.txt': 'AT pattern2.txt', '/root_path/US/rest_of_path/pattern1.txt': 'US pattern1.txt', '/root_path/US/rest_of_path/pattern2.txt': 'US pattern2.txt'} mock_get_files_from_sftp.return_value = response mock_task_status.get_values.return_value = [ 'AT pattern1.txt', 'AT pattern2.txt', 'US pattern1.txt', 'US pattern2.txt'] result = ftp_amazon_marketshare_tasks.check_for_new_files( activity=MagicMock(), feed_name='feed_name', date='2016-01-01', ftp_creds=sftp_conf_mock, ftp_path_template='/root_path/{country_code}/rest_of_path', file_pattern='pattern') mock_delete_status.assert_not_called() assert result == {'stop': True} @patch('feed_ingestion.tasks.ftp_amazon_marketshare_tasks.garcon_ftp.' 'copy_file_from_ftp_to_s3') def test_fetch_from_drop_location( mock_copy_file_from_ftp_to_s3, sftp_conf_mock): """Test test_fetch_from_drop_location method.""" def sf_copy_file( activity, ftp_creds, ftp_path, file_name, s3_archive_path, new_file_name): """copy_file_from_ftp_to_s3 side effect function.""" if file_name == 'filename1': return {'status': False} return { 'file_size': 1, 'status': True} files_on_ftp = { 'account_name1': { '/root_path/AT/rest_of_path/filename1': 'AT filename1', '/root_path/US/rest_of_path/filename2': 'US filename2', '/root_path/US/rest_of_path/filename3': 'DE filename5', '/root_path/US/rest_of_path/filename4': 'DE filename5_Rev' }, 'account_name2': { '/root_path/AT/rest_of_path/filename5': 'AT filename3', '/root_path/US/rest_of_path/filename6': 'US filename4', '/root_path/US/rest_of_path/filename7': 'DE filename6', '/root_path/US/rest_of_path/filename8': 'DE filename6_Rev', '/root_path/AT/rest_of_path/filename9': 'US filename9', '/root_path/AT/rest_of_path/filename10': 'US filename9_revised' }, } mock_copy_file_from_ftp_to_s3.side_effect = sf_copy_file result = ftp_amazon_marketshare_tasks.fetch_from_drop_location( activity=MagicMock(), date='2016-01-01', ftp_creds=sftp_conf_mock, files_on_ftp=files_on_ftp, s3_archive_path='s3://test_path/' )['source_files_dict']['files'] expected_response = [ { 'file_name': 'AT filename1', 'file_size': -1, 'found': False }, { 'file_name': 'US filename2', 'file_size': 1, 'found': True }, { 'file_name': 'DE filename5_Rev', 'file_size': 1, 'found': True }, { 'file_name': 'AT filename3', 'file_size': 1, 'found': True }, { 'file_name': 'US filename4', 'file_size': 1, 'found': True }, { 'file_name': 'DE filename6_Rev', 'file_size': 1, 'found': True }, { 'file_name': 'US filename9_revised', 'file_size': 1, 'found': True } ] assert sorted(result, key=lambda d: d['file_name']) == \ sorted(expected_response, key=lambda d: d['file_name'])