"""Unit tests for S3 related utils and helpers of Feed Ingestion workflows.""" from io import StringIO import os import re from unittest.mock import MagicMock, patch import boto3 from botocore.exceptions import ClientError from garcon_contrib.aws.utils import garcon_s3 from moto import mock_aws import pytest from feed_ingestion.flows import helpers from feed_ingestion.util.aws import s3 @pytest.fixture def s3_client(): with mock_aws(): yield boto3.client('s3') @pytest.fixture def s3_bucket(s3_client): bucket = 'abucket' s3_client.create_bucket(Bucket=bucket) yield bucket def test_convert_zip_to_gzip_on_s3(monkeypatch): """Testing the S3 functions called by convert_zip_to_gzip_on_s3.""" monkeypatch.setattr( os.path, 'basename', MagicMock(return_value='file_name')) monkeypatch.setattr(os, 'remove', MagicMock(return_value=None)) monkeypatch.setattr(helpers, 'download', MagicMock(return_value=None)) monkeypatch.setattr( s3, '_convert_zip_to_gzip', MagicMock(return_value=None)) monkeypatch.setattr( helpers, 'upload_raw_file_to_s3', MagicMock(return_value=None)) resp = s3.convert_zip_to_gzip_on_s3( MagicMock(), 'zip_path', 'gzip_path', './') assert resp.get('converted_file_s3_full_path') == 'gzip_path' helpers.upload_raw_file_to_s3.assert_any_call( './file_name', 'gzip_path') s3._convert_zip_to_gzip.assert_any_call( './file_name', './file_name', './', False) assert s3._convert_zip_to_gzip.call_count == 1 assert os.remove.called s3.convert_zip_to_gzip_on_s3( MagicMock(), 'zip_path', 'gzip_path', './', extract_original_filename=True) assert s3._convert_zip_to_gzip.call_count == 2 s3._convert_zip_to_gzip.assert_any_call( './file_name', './file_name', './', True) s3.convert_zip_to_gzip_on_s3( MagicMock(), 'zip_path', 'gzip_path', './', extract_original_filename=True) def test_copy_s3_key(s3_client, s3_bucket): """Testing the s3 functions called by copy_s3_key.""" old_path = 'test_old' new_path = 'test_new' s3_client.put_object(Bucket=s3_bucket, Key=old_path, Body='123') old_path_full = f's3://{s3_bucket}/{old_path}' new_path_full = f's3://{s3_bucket}/{new_path}' s3.copy_s3_key(old_path_full, new_path_full) s3_client.head_object(Bucket=s3_bucket, Key=old_path) s3_client.head_object(Bucket=s3_bucket, Key=new_path) def test_get_key_size(s3_client, s3_bucket): """Testing the S3 functions called by get_key_size.""" input_path = f's3://{s3_bucket}/path/file' s3_client.put_object(Bucket=s3_bucket, Key='path/file', Body='12345') result = s3.get_key_size(input_path) assert result == 5 def test_get_key_size_expected_bucket_owner_env(monkeypatch, s3_bucket): input_path = f's3://{s3_bucket}/path/file' mock_boto3 = MagicMock() mock_boto3.client.return_value.get_object.return_value = { 'ContentLength': 5} monkeypatch.setattr(s3, 'boto3', mock_boto3) monkeypatch.setattr(s3.config, 'EXPECTED_BUCKET_OWNER', 'env-owner') result = s3.get_key_size(input_path) assert result == 5 mock_boto3.client.return_value.get_object.assert_called_once_with( Bucket=s3_bucket, Key='path/file', ExpectedBucketOwner='env-owner' ) def test_get_key_size_expected_bucket_owner_provided(monkeypatch, s3_bucket): input_path = f's3://{s3_bucket}/path/file' mock_boto3 = MagicMock() mock_boto3.client.return_value.get_object.return_value = { 'ContentLength': 5} monkeypatch.setattr(s3, 'boto3', mock_boto3) result = s3.get_key_size( input_path, expected_bucket_owner='provided-owner') assert result == 5 mock_boto3.client.return_value.get_object.assert_called_once_with( Bucket=s3_bucket, Key='path/file', ExpectedBucketOwner='provided-owner' ) @mock_aws @pytest.mark.skip('Very badly written test. But tested method is not used') # TODO: remove tested method and this test def test_expand_s3_csv_with_columns( monkeypatch, tmp_path, s3_client, s3_bucket): """Test transformation of CSV files. (And copying them to separate directory on S3). """ file_to_transform = ( 's3://bucket/ProperStock/archives/2016-09-29/' 'GoodsIn_ESSN_2016-09-29.csv') s3_bucket, bucket_path = garcon_s3.extract_bucket_path( file_to_transform) content = 'test1, test2' + '\n' + 'test3, test4' s3_client.put_object( Bucket=s3_bucket, Key=bucket_path, Body=content) columns_header = ('file_date', 'file_name', 'ingestion_timestamp') columns_values = ('2016-09-29', file_to_transform, '2016-10-07T13:26:11') path_to_unload = '/ProperStock/snowflake/2016-09-29' s3.expand_s3_csv_with_columns( file_to_transform, columns_header, columns_values, path_to_unload) # download content of the file to string result_file = tmp_path / 'result.csv' s3_client.download_file( s3_bucket, bucket_path, result_file) str_content = result_file.read_text() assert (str_content.split('\r\n')[0] == 'test1, test2,file_date,file_name,ingestion_timestamp') assert (','.join(str_content.split( '\r\n')[1].split(',')[:-1]).rstrip(',') == 'test3, test4,{date},{file}'.format( date=file_to_transform.split('/')[-2], file=file_to_transform)) assert re.match( r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}', str_content.split('\r\n')[1].split(',')[-1]) def test_delete_file(s3_client, s3_bucket): """Test delete file helper.""" key = 'path/file' full_path = f's3://{s3_bucket}/{key}' s3_client.put_object(Bucket=s3_bucket, Key=key, Body='12') s3_client.head_object(Bucket=s3_bucket, Key=key) s3.delete_file(full_path) with pytest.raises(ClientError): s3_client.head_object(Bucket=s3_bucket, Key=key) def test_get_file_size(s3_client, s3_bucket): """Test get file size.""" s3_client.put_object( Bucket=s3_bucket, Key='some_path', Body='1' * 1024) s3_full_path = f's3://{s3_bucket}/some_path' size = s3.get_file_size(s3_full_path) assert size == 1024 def test_upload_processed_to_s3(s3_client, s3_bucket): """Test upload_processed_to_s3 task.""" key = 'path/test_filename.tsv.gz' result = s3.upload_processed_to_s3( StringIO(initial_value='1234'), f's3://{s3_bucket}/{key}') assert result == {'source_files_dict': { 'files': [{ 'file_name': 'test_filename.tsv.gz', 'found': True, 'file_size': 4}]}} assert s3_client.head_object( Bucket=s3_bucket, Key=key)['ContentLength'] >= 0 @mock_aws @patch('feed_ingestion.util.aws.s3.smart_open.smart_open') def test_get_source_files_content(mock_smart_open): """Test get_source_files_content.""" s3_path = 's3://bucket/path/' source_files_dict = {'files': [ {'file_name': 'test1'}, {'file_name': 'test2'} ]} list(s3.get_source_files_content(s3_path, source_files_dict)) assert mock_smart_open.called assert mock_smart_open().__enter__().__iter__.called @mock_aws def test_upload_on_s3(monkeypatch): """Test upload_on_s3.""" client = MagicMock() monkeypatch.setattr( boto3, 'client', MagicMock(return_value=client)) fd = MagicMock() fd.name = 'test.csv' s3.upload_on_s3('bucket', 'test/path/', 'filename', fd) client.upload_file.assert_called_with( 'test.csv', 'bucket', 'test/path/filename', ExtraArgs={'ExpectedBucketOwner': '437795906767'})