import gzip import json import unittest from datetime import datetime import pytest from requests import Response from slz_appreciationengine_scrapper.const import RATE_LIMIT_HOUR_HEADER from slz_appreciationengine_scrapper.content_status_service import ContentStatusService from slz_appreciationengine_scrapper.dsp.appreciationengine.clients import PaginateClientPartial from slz_appreciationengine_scrapper.dsp.entities import MetaDataKey from slz_appreciationengine_scrapper.entities import Job def resp_mock_activityfeed_us_columbia_partial(): data = [ { 'totalSize': 1000, 'items': [{ 'RecordDate': '2020-11-04 23:17:16' }, ] }, # request to check filtering rule { 'totalSize': 1000, 'items': [{ 'RecordDate': '2020-11-04 23:17:16' }, ] }, { 'totalSize': 1000, 'items': [ { 'RecordDate': '2020-11-04 23:17:17' }, { 'RecordDate': '2020-11-04 23:17:18' }, { 'RecordDate': '2020-11-04 23:17:19' }, ] }, { 'totalSize': 1000, 'items': [ { 'RecordDate': '2020-11-04 23:17:20' }, { 'RecordDate': '2020-11-04 23:17:21' }, { 'RecordDate': '2020-11-04 23:17:22' }, { 'RecordDate': '2020-11-04 23:17:23' }, { 'RecordDate': '2020-11-04 23:17:25' }, ] }, { 'totalSize': 1, 'items': [{ 'RecordDate': '2020-11-04 23:17:27' }] } ] rate_limit = 7 for dt in data: resp_mock = unittest.mock.MagicMock(spec=Response) resp_mock.status_code = 200 resp_mock.ok = True resp_mock.url = '' resp_mock.content = bytes(json.dumps(dt), encoding='utf8') resp_mock.headers = {RATE_LIMIT_HOUR_HEADER: rate_limit} rate_limit -= 1 yield resp_mock @pytest.mark.integration @unittest.mock.patch( 'slz_appreciationengine_scrapper.dsp.appreciationengine.ae_io.time.sleep', return_value=None ) @unittest.mock.patch('slz_appreciationengine_scrapper.dsp.appreciationengine.clients.get_secret') @unittest.mock.patch('requests.get') @unittest.mock.patch( 'slz_appreciationengine_scrapper.dsp.appreciationengine.clients.datetime', wraps=datetime, now=lambda *args, **kwargs: datetime(year=2021, month=6, day=5, hour=1, minute=1, second=1) ) def test_downloading_paginate_partial( datetime_mock, requests_get, get_secret, patched_time_sleep, db, params, s3_client, ae_mapping, ): for bucket in ['bucket-archive-quarantine', 'bucket-decompressed-quarantine']: s3_client.create_bucket(Bucket=bucket) uow_dict = { 'uow_id': 'appreciationengine-20201104-sme-activityfeed-v1', 'dsp': 'appreciationengine', 'report_type': 'activityfeed', 'version': 'v1', 'report_date': '2020-11-04', 'licensor': 'sme', 'extension': 'json', 'context': 'US_Columbia', } # Client initialization logger = unittest.mock.Mock() timer = unittest.mock.Mock() client = PaginateClientPartial(logger=logger) # Mock secrets get_secret.return_value = json.dumps({ "Sony Music US - Columbia": "afa0d1d", }) params.scrapper.min_rate_limit_remaining_hour = 5 client.configure(params) job = Job.from_dict(uow_dict) expected_file_content = ''.join( [ '{"RecordDate": "2020-11-04 23:17:16"}\n', '{"RecordDate": "2020-11-04 23:17:17"}\n', '{"RecordDate": "2020-11-04 23:17:18"}\n', '{"RecordDate": "2020-11-04 23:17:19"}\n', '{"RecordDate": "2020-11-04 23:17:20"}\n', '{"RecordDate": "2020-11-04 23:17:21"}\n', '{"RecordDate": "2020-11-04 23:17:22"}\n', '{"RecordDate": "2020-11-04 23:17:23"}\n', ] ) # Mock response from source API requests_get.return_value.__enter__.side_effect = resp_mock_activityfeed_us_columbia_partial() content_status_service = ContentStatusService(logger=logger, db_conn=db) meta, meta_data, io_exception = client.download( job, chunk_size=1, timer=timer, timeslot=None, content_status_service=content_status_service ) # read data from mocked S3 buckets path = 'appreciationengine/activityfeed/v1/report_date=2020-11-04/report_licensor=sme' actual_decompressed = s3_client.get_object( Bucket='bucket-decompressed-quarantine', Key=f'{path}/US_Columbia_20201104_20210605010101.json', )['Body'].read() actual_compressed = s3_client.get_object( Bucket='bucket-archive-quarantine', Key=f'{path}/US_Columbia_20201104_20210605010101.json.gz', )['Body'].read() assert expected_file_content == actual_decompressed.decode('utf8') assert expected_file_content == gzip.decompress(actual_compressed).decode('utf8') assert meta_data[MetaDataKey.NEW_START_DATE.value] == "2020-11-04 23:17:24" def resp_mock_activityfeed_us_columbia_partial_small_resp(): data = [ { 'totalSize': 1, 'items': [{ 'RecordDate': '2020-11-04 19:31:16' }, ] }, # request to check filtering rule { 'totalSize': 1, 'items': [{ 'RecordDate': '2020-11-04 19:31:16' }, ] }, { 'totalSize': 0, 'items': [] } ] rate_limit = 7 for dt in data: resp_mock = unittest.mock.MagicMock(spec=Response) resp_mock.status_code = 200 resp_mock.ok = True resp_mock.url = '' resp_mock.content = bytes(json.dumps(dt), encoding='utf8') resp_mock.headers = {RATE_LIMIT_HOUR_HEADER: rate_limit} rate_limit -= 1 yield resp_mock @pytest.mark.integration @unittest.mock.patch( 'slz_appreciationengine_scrapper.dsp.appreciationengine.ae_io.time.sleep', return_value=None ) @unittest.mock.patch('slz_appreciationengine_scrapper.dsp.appreciationengine.clients.get_secret') @unittest.mock.patch('requests.get') @unittest.mock.patch( 'slz_appreciationengine_scrapper.dsp.appreciationengine.clients.datetime', wraps=datetime, now=lambda *args, **kwargs: datetime(year=2021, month=6, day=5, hour=1, minute=1, second=1) ) def test_downloading_paginate_partial_small_resp( datetime_mock, requests_get, get_secret, patched_time_sleep, db, params, s3_client, ae_mapping, logger_test ): for bucket in ['bucket-archive-quarantine', 'bucket-decompressed-quarantine']: s3_client.create_bucket(Bucket=bucket) uow_dict = { 'uow_id': 'appreciationengine-20201104-sme-activityfeed-v1', 'dsp': 'appreciationengine', 'report_type': 'activityfeed', 'version': 'v1', 'report_date': '2020-11-04', 'licensor': 'sme', 'extension': 'json', 'context': 'US_Columbia', } # Client initialization logger = unittest.mock.Mock() timer = unittest.mock.Mock() client = PaginateClientPartial(logger=logger_test) # Mock secrets get_secret.return_value = json.dumps({ "Sony Music US - Columbia": "afa0d1d", }) params.scrapper.min_rate_limit_remaining_hour = 5 client.configure(params) job = Job.from_dict(uow_dict) expected_file_content = ''.join([ '{"RecordDate": "2020-11-04 19:31:16"}\n', ]) # Mock response from source API requests_get.return_value.__enter__.side_effect = resp_mock_activityfeed_us_columbia_partial_small_resp( ) meta_data_initial = {MetaDataKey.NEW_START_DATE.value: "2020-11-04 19:30:01"} content_status_service = ContentStatusService(logger=logger, db_conn=db) meta, meta_data, io_exception = client.download( job, chunk_size=1, timer=timer, meta_data=meta_data_initial, timeslot=None, content_status_service=content_status_service ) # read data from mocked S3 buckets path = 'appreciationengine/activityfeed/v1/report_date=2020-11-04/report_licensor=sme' actual_decompressed = s3_client.get_object( Bucket='bucket-decompressed-quarantine', Key=f'{path}/US_Columbia_20201104_20210605010101.json', )['Body'].read() actual_compressed = s3_client.get_object( Bucket='bucket-archive-quarantine', Key=f'{path}/US_Columbia_20201104_20210605010101.json.gz', )['Body'].read() assert expected_file_content == actual_decompressed.decode('utf8') assert expected_file_content == gzip.decompress(actual_compressed).decode('utf8') assert meta_data[MetaDataKey.PARTIAL_DOWNLOAD_COMPLETED_AT.value] is not None