import json import os import unittest from datetime import datetime from unittest.mock import MagicMock, Mock import boto3 import smart_open from db_schema.schemas.exploration import DisassembleContentStatus, DisassembleStatusEnum from moto import mock_s3 from parameterized import parameterized from exp_sf_runner_lambda.entities import LambdaConfig from exp_sf_runner_lambda.exceptions import DbxConfigError from exp_sf_runner_lambda.handler import get_config_from_file from exp_sf_runner_lambda.service import SfService FIXTURES_PATH = os.path.join(os.path.dirname(__file__), 'fixtures') class MockContextManager: def __enter__(self): return self def __exit__(self, exc_type, exc, tb): pass @mock_s3 class SfServiceTestCase(unittest.TestCase): def setUp(self): boto3.setup_default_session() self.conn = boto3.resource('s3', region_name='us-east-1') self.bucket = 'dev-test-bucket' self.conn.create_bucket(Bucket=self.bucket) expected = 'Some bytes data' self.data = bytes(expected, 'utf-8') self.s3_path = f's3://{self.bucket}/initial_test_file.txt' with smart_open.open(self.s3_path, 'wb') as fout: fout.write(self.data) self.sf_client = Mock() self._pg_repo = Mock() self._pg_repo.advisory_locked_transaction = MagicMock(MockContextManager()) self._dbx_config = get_config_from_file( os.path.join(FIXTURES_PATH, 'exploration-etl-config.json') ) lambda_cfg = LambdaConfig( sf_arn='SF_ARN', max_active_runs_dbx=2, max_active_runs_total=2, dest_bucket='DEST_BUCKET', exp_config_path='EXP_CONFIG_PATH', max_filesize_not_split=500, retry_failed_limit=5, trace_id='trace_id', optional_config={'RDS_SECRET_KEY': 'test/key'} ) self.service = SfService( logger=Mock(), config=lambda_cfg, dbx_configs=self._dbx_config, sf_client=self.sf_client, pg_repo=self._pg_repo, sf_repo=Mock() ) self.service._s3_client = self.conn @parameterized.expand( [ ( 'apple-20190903-sme-amContent-v1_2', 123, 'apple-20190903-sme-amContent-v1_2_123_2020-01-15T01.02.03' ), ( 'gras-20200507-sme-dim_rec_project_product_version-v1-rp20190906T105504', 123123123, 'gras-20200507-sme-dim_rec_project_product_version-_123123123_2020-01-15T01.02.03' ), ] ) @unittest.mock.patch('exp_sf_runner_lambda.service.datetime') def test_service__generate_sf_name(self, uow_id, dcs_id, expected, dt): dt.today.return_value = datetime(2020, 1, 15, 1, 2, 3, 4) got = self.service._generate_sf_name(uow_id, dcs_id) self.assertEqual(got, expected) def test_get_filesize_ok(self): actual = self.service._get_filesize(self.s3_path) expected = len(self.data) self.assertEqual(actual, expected) @parameterized.expand( [ ('amazonmusicunlimited-20191216-theorchard-activity-v1', 'default'), ('amazonprime-20191216-theorchard-activity-v1', 'amazonprime') ] ) def test_get_dbx_params_correct(self, uow_id, expected_from_conf): actual = self.service._get_dbx_params(uow_id) expected = self._dbx_config[expected_from_conf] self.assertEqual(actual, expected) def test_get_dbx_params_raise_exc(self): self.service._dbx_configs = {} self.assertRaises(DbxConfigError, self.service._get_dbx_params, "uow") @unittest.mock.patch('exp_sf_runner_lambda.service.datetime') def test__start_sf_without_dbx_usual_flow(self, dt): dt.today.return_value = datetime(2020, 1, 15, 1, 2, 3, 4) rec = DisassembleContentStatus( disassemble_content_status_id=1, content_name='initial_test_file.txt', payload={ 'uow_id': 'amazonprime-20191216-theorchard-activity-v1', 'source_path': self.s3_path, 'context': 'users', } ) sf_name, use_dbx = self.service._start_sf(rec) self.assertEqual(use_dbx, False) name = 'amazonprime-20191216-theorchard-activity-v1_1_2020-01-15T01.02.03' self.sf_client.start_execution.assert_called_with( name=name, stateMachineArn='SF_ARN', input=json.dumps( { 'name': name, 'uow_id': 'amazonprime-20191216-theorchard-activity-v1', 'context': 'users', 'disassemble_content_status_id': 1, 'content_name': 'initial_test_file.txt', 'source_path': self.s3_path, 'dest_bucket': 'DEST_BUCKET', 'action': 'ARCHIVE_COPY', 'trace_id': 'trace_id', 'optional_config': { 'RDS_SECRET_KEY': 'test/key' }, 'archive_params': { 'drop_from_head': 3, 'drop_from_tail': 2 } } ), ) @unittest.mock.patch('exp_sf_runner_lambda.service.datetime') def test__start_sf_without_dbx_parquets(self, dt): source_path = 's3://dev-archive/path1c000.snappy.parquet,s3://dev-archive/path2c000.snappy.parquet' dt.today.return_value = datetime(2020, 1, 15, 1, 2, 3, 4) rec = DisassembleContentStatus( disassemble_content_status_id=1, content_name='initial_test_file.txt', payload={ 'uow_id': 'youtubereporting-20201108-sme-content_owner_asset-a2', 'source_path': source_path, 'context': 'content_owner_asset', 'action': 'COPY' } ) sf_name, use_dbx = self.service._start_sf(rec) self.assertEqual(use_dbx, False) name = 'youtubereporting-20201108-sme-content_owner_asset-a2_1_2020-01-15T01.02.03' self.sf_client.start_execution.assert_called_with( name=name, stateMachineArn='SF_ARN', input=json.dumps( { 'name': name, 'uow_id': 'youtubereporting-20201108-sme-content_owner_asset-a2', 'context': 'content_owner_asset', 'disassemble_content_status_id': 1, 'content_name': 'initial_test_file.txt', 'source_path': source_path, 'dest_bucket': 'DEST_BUCKET', 'action': 'COPY', 'trace_id': 'trace_id', 'optional_config': { 'RDS_SECRET_KEY': 'test/key' }, 'archive_params': { 'drop_from_head': 0, 'drop_from_tail': 0 } } ), ) @unittest.mock.patch('exp_sf_runner_lambda.service.datetime') def test__start_sf_with_dbx(self, dt): dt.today.return_value = datetime(2020, 1, 15, 1, 2, 3, 4) fsize = 600 * 1024**2 self.service._get_filesize = lambda x: fsize rec = DisassembleContentStatus( disassemble_content_status_id=1, content_name='initial_test_file.txt', payload={ 'uow_id': 'amazonprime-20191216-theorchard-activity-v1', 'source_path': self.s3_path, 'context': 'users' } ) sf_name, use_dbx = self.service._start_sf(rec) self.assertEqual(use_dbx, True) name = 'amazonprime-20191216-theorchard-activity-v1_1_2020-01-15T01.02.03' self.sf_client.start_execution.assert_called_with( name=name, stateMachineArn='SF_ARN', input=json.dumps( { 'name': name, 'uow_id': 'amazonprime-20191216-theorchard-activity-v1', 'context': 'users', 'disassemble_content_status_id': 1, 'content_name': 'initial_test_file.txt', 'source_path': self.s3_path, 'dest_bucket': 'DEST_BUCKET', 'action': 'SPLIT_ARCHIVE_COPY_DBX', 'trace_id': 'trace_id', 'optional_config': { 'RDS_SECRET_KEY': 'test/key' }, 'databricks_job_id': 260, 'databricks_spark_params': ['--class1'], } ), ) def test_start_sf_correct(self): def start_sf_mock(x): return 'test_sf_name', True self.service._start_sf = start_sf_mock rec = DisassembleContentStatus(disassemble_content_status_id=1, ) result, use_dbx = self.service.start_sf(rec) self.assertEqual(result, True) self.assertEqual(use_dbx, True) self.service._pg_repo.update_disassemble_status.assert_called_once_with( disassemble_cs_id=1, process_status=DisassembleStatusEnum.IN_PROGRESS, params={'sf_execution_name': 'test_sf_name'} ) def test_start_sf_correctHandling_exc(self): def start_sf_mock(x): raise Exception self.service._start_sf = start_sf_mock rec = DisassembleContentStatus(disassemble_content_status_id=1, ) result, use_dbx = self.service.start_sf(rec) self.assertEqual(result, False) self.assertEqual(use_dbx, False) self.service._pg_repo.update_disassemble_status.assert_called_once_with( disassemble_cs_id=1, process_status=DisassembleStatusEnum.FAILED, ) def test_retry_if_failed_no_retries(self): self.service._config.retry_failed_limit = 0 rec = DisassembleContentStatus( disassemble_content_status_id=1, disassemble_status=DisassembleStatusEnum.FAILED ) self.service._retry_if_failed(rec) self._pg_repo.update_disassemble_status.assert_called_once_with( disassemble_cs_id=1, process_status=DisassembleStatusEnum.FAILED, is_failure_transient=False ) self._pg_repo.create_disassemble_content_status.assert_not_called() def test_retry_if_failed_has_retries(self): self.service._config.retry_failed_limit = 1 rec = DisassembleContentStatus( disassemble_content_status_id=1, disassemble_status=DisassembleStatusEnum.FAILED, payload={ 'uow_id': 'amazonprime-20191216-theorchard-activity-v1', 'source_path': self.s3_path, 'context': 'users' } ) # should create new dcs self.service._retry_if_failed(rec) self._pg_repo.create_disassemble_content_status.assert_called() self._pg_repo.update_disassemble_status.assert_called_once_with( disassemble_cs_id=1, process_status=DisassembleStatusEnum.REPLACED, ) def test_retry_if_failed_limit_reached_after_several_retries(self): rec = DisassembleContentStatus( disassemble_content_status_id=1, disassemble_status=DisassembleStatusEnum.FAILED, retried_disassemble_content_status_id=123, ) self._pg_repo.get_retries_count.side_effect = [ 5, ] self.service._retry_if_failed(rec) self._pg_repo.update_disassemble_status.assert_called_once_with( disassemble_cs_id=1, process_status=DisassembleStatusEnum.FAILED, is_failure_transient=False )