"""Unit tests for Amazon DataPulse tasks.""" from unittest.mock import MagicMock, patch import pytest from feed_ingestion.flows.amazon_datapulse import tasks _date = '2024-01-15' @pytest.fixture def mock_get_overall_status(): """Yield get_overall_status mock.""" with patch.object(tasks.garcon_feed_status, 'get_overall_status') as mock_obj: yield mock_obj @pytest.fixture def mock_set_overall_status(): """Yield set_overall_status mock.""" with patch.object(tasks.garcon_feed_status, 'set_overall_status') as mock_obj: yield mock_obj @pytest.fixture def mock_delete_status(): """Yield delete_status mock.""" with patch.object(tasks.garcon_feed_status, 'delete_status') as mock_obj: yield mock_obj @pytest.fixture def mock_get_staging_raw_database_schema(): """Yield delete_status mock.""" with patch.object(tasks.config, 'get_staging_raw_database_schema') as mock_obj: mock_obj.side_effect = ( lambda report: (f'test_db_{report}', f'test_schema_{report}', f'staging_raw_amazon_{report}') ) yield mock_obj @pytest.fixture def mock_athena(): """Yield athena module mock.""" with patch.object(tasks, 'athena') as mock_obj: yield mock_obj @pytest.fixture def mock_s3_tasks(): """Yield s3_tasks module mock.""" with patch.object(tasks, 's3_tasks') as mock_obj: mock_obj.source_files.return_value = { 'source_files_dict': {'files': [{'file_name': 'part-0.parquet'}]} } yield mock_obj @pytest.fixture def expected_bootstrap_response(): """Return expected response from bootstrap task.""" return { 'feed_name': 'amazon_datapulse_fraud_report', 'executor_feed_name': 'amazon_datapulse', 'report_name': 'fraud_report', 'date': '2024-01-15', 'secrets_path': 'amazon_datapulse', 'archive_bucket': 'dev-cucumbers', 'archive_path': 'AmazonDataPulse/fraud_report/full/2024-01-15/', 's3_dir_path': ( 's3://dev-cucumbers/AmazonDataPulse/fraud_report/full/2024-01-15/' ), 'athena_source_table': 'fraud_report', 'staging_raw_table': 'staging_raw_amazon_fraud_report', 'sfdb_params': {}, 'kwargs': {'partition': {}, 'staging_raw_database': 'test_db_fraud_report', 'staging_raw_schema': 'test_schema_fraud_report'}, 'partition': {}, } class TestBootstrap: """Tests for the bootstrap task.""" class TestFraudReport: """Tests for fraud_report.""" report_name = 'fraud_report' def test_bootstrap_normal( self, expected_bootstrap_response, mock_get_overall_status, mock_delete_status, mock_get_staging_raw_database_schema ): """Test bootstrap with normal execution returns full context.""" result = tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name ) assert result == expected_bootstrap_response mock_delete_status.assert_not_called() def test_bootstrap_reload( self, expected_bootstrap_response, mock_get_overall_status, mock_delete_status, mock_get_staging_raw_database_schema, ): """Test bootstrap with reload=True.""" result = tasks.bootstrap( MagicMock(), _date, reload='True', report_name=self.report_name ) assert result == expected_bootstrap_response mock_delete_status.assert_called_once_with( 'amazon_datapulse_fraud_report', '2024-01-15' ) def test_bootstrap_inconsistent_dates( self, expected_bootstrap_response, mock_get_overall_status, mock_delete_status, ): partition = { 'cadence': 'daily', 'region': 'us-east-1', 'report_date': '2026-04-02', 'entity_name': 'Sony', 'service': 'UNLIMITED', 'music_territory': 'MX', } with pytest.raises(ValueError, match='Inconsistent dates'): tasks.bootstrap( MagicMock(), '2026-04-03', reload='True', report_name=self.report_name, partition=partition, ) mock_delete_status.assert_not_called() def test_bootstrap_already_ingested(self, mock_get_overall_status): """Test bootstrap returns stop when workflow already ingested.""" mock_get_overall_status.return_value = 'INGESTED' result = tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name ) assert result == { 'stop': True, 'message': ( 'amazon_datapulse_fraud_report ' 'is already ingested for 2024-01-15' ), } def test_bootstrap_with_partition( self, expected_bootstrap_response, mock_get_overall_status, mock_delete_status, mock_get_staging_raw_database_schema, ): """Test bootstrap with valid partition adds it to result.""" partition = { 'cadence': 'daily', 'region': 'us-east-1', 'report_date': '2024-01-15', 'entity_name': 'Sony', 'service': 'UNLIMITED', 'music_territory': 'MX', } result = tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition=partition, ) s3_path = ( 'AmazonDataPulse/fraud_report/partitioned' '/cadence=daily/region=us-east-1/report_date=2024-01-15' '/entity_name=Sony/service=UNLIMITED/music_territory=MX/' ) expected_bootstrap_response = { **expected_bootstrap_response, 'kwargs': { 'partition': partition, 'staging_raw_database': 'test_db_fraud_report', 'staging_raw_schema': 'test_schema_fraud_report' }, 'archive_path': s3_path, 's3_dir_path': f's3://dev-cucumbers/{s3_path}', 'feed_name': 'amazon_datapulse_fraud_report_daily' '_us_east_1_2024_01_15_sony_unlimited_mx', } assert result == { **expected_bootstrap_response, 'partition': partition, } def test_bootstrap_invalid_partition_key( self, mock_get_overall_status ): """Test bootstrap raises ValueError for unknown partition key.""" regex = ( 'Invalid partition "\\[\'unknown_key\'\\]" ' 'for report "fraud_report"' ) with pytest.raises(ValueError, match=regex): tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition={'unknown_key': 'value'}, ) def test_bootstrap_invalid_partition_value_type( self, mock_get_overall_status ): """Test bootstrap when value type doesn't match.""" partition = { 'cadence': 'daily', 'region': 123, 'report_date': '2026-04-02', 'entity_name': 'Sony', 'service': 'UNLIMITED', 'music_territory': 'MX', } regex = 'Invalid type for partition key "region"' with pytest.raises(ValueError, match=regex): tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition=partition, ) class TestDailyPlayEvents: """Tests for daily_play_events report (year/month/day partitions).""" report_name = 'daily_play_events' def test_bootstrap_normal( self, mock_get_overall_status, mock_delete_status ): """Test bootstrap with no partition returns full archive path.""" result = tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name ) assert result['feed_name'] == ( 'amazon_datapulse_daily_play_events' ) assert result['report_name'] == self.report_name assert result['archive_path'] == ( 'AmazonDataPulse/daily_play_events/full/2024-01-15/' ) assert result['athena_source_table'] == 'daily_play_events' assert result['staging_raw_table'] == ( 'staging_raw_amazon_daily_play_events' ) assert result['partition'] == {} def test_bootstrap_with_partition( self, mock_get_overall_status, mock_delete_status ): """Test bootstrap with year/month/day partition.""" partition = { 'service': 'UNLIMITED', 'region': 'NA', 'marketplace': 'ATVPDKIKX0DER', 'territory': 'US', 'year': '2024', 'month': '01', 'day': '15', } result = tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition=partition, ) assert result['feed_name'] == ( 'amazon_datapulse_daily_play_events_unlimited' '_na_atvpdkikx0der_us_2024_01_15' ) assert result['archive_path'] == ( 'AmazonDataPulse/daily_play_events/partitioned' '/service=UNLIMITED/region=NA' '/marketplace=ATVPDKIKX0DER/territory=US' '/year=2024/month=01/day=15/' ) assert result['partition'] == partition def test_bootstrap_inconsistent_partition_date( self, mock_get_overall_status ): """Test bootstrap raises when year/month/day != context_date.""" partition = { 'service': 'UNLIMITED', 'region': 'NA', 'marketplace': 'ATVPDKIKX0DER', 'territory': 'US', 'year': '2024', 'month': '02', 'day': '01', } with pytest.raises(ValueError, match='Inconsistent dates'): tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition=partition, ) def test_bootstrap_invalid_partition_key( self, mock_get_overall_status ): """Test bootstrap raises when partition keys don't match config.""" with pytest.raises(ValueError, match='Invalid partition'): tasks.bootstrap( MagicMock(), _date, reload=None, report_name=self.report_name, partition={'unknown_key': 'value'}, ) @pytest.fixture def mock_task_status(monkeypatch): """Yield task_status mocks to bypass DynamoDB in check_status decorator.""" from feed_ingestion.util import task_status monkeypatch.setattr( task_status, 'is_completed_task', MagicMock(return_value=False) ) monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock()) class TestFetchFromAthena: """Tests for the fetch_from_athena task.""" s3_path = 'amazon_datapulse/fraud_report/2024-01-15/theorchard/' feed_name = 'amazon_datapulse_fraud_report_theorchard' bucket = 'some_bucket' report_name = 'fraud_report' def test_fetch_from_athena_success( self, mock_task_status, mock_set_overall_status, mock_athena, mock_s3_tasks, ): """Test fetch_from_athena marks files found and sets DOWNLOADED.""" result = tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=self.s3_path, ) assert result == mock_s3_tasks.source_files.return_value for file in result['source_files_dict']['files']: assert file['found'] is True mock_set_overall_status.assert_called_once_with( self.feed_name, _date, 'DOWNLOADED' ) def test_fetch_from_athena_runs_athena_query( self, mock_task_status, mock_set_overall_status, mock_athena, mock_s3_tasks, snapshot, ): """Test fetch_from_athena calls athena.run_query with UNLOAD.""" tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=self.s3_path, ) mock_athena.run_query.assert_called_once() sql_query = mock_athena.run_query.call_args_list[0][1]['athena_query'] snapshot.assert_match(sql_query, 'athena_query.sql') call_kwargs = mock_athena.run_query.call_args.kwargs assert call_kwargs['destination_s3_bucket'] == self.bucket assert call_kwargs['destination_s3_path'] == self.s3_path assert call_kwargs['use_unload_query'] is True def test_fetch_from_athena_empty_dataset( self, mock_task_status, mock_set_overall_status, mock_athena, mock_s3_tasks, ): """Test fetch_from_athena raises ValueError on empty result set.""" mock_s3_tasks.source_files.side_effect = ValueError with pytest.raises(ValueError, match='Empty dataset'): tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=self.s3_path, ) mock_set_overall_status.assert_called_once_with( self.feed_name, _date, 'NOT_AVAILABLE' ) def test_fetch_from_athena_invalid_s3_path( self, mock_task_status, mock_athena, mock_s3_tasks ): """Test fetch_from_athena raises ValueError for missing trailing /.""" with pytest.raises( ValueError, match='destination_s3_path should end with /' ): tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=( 'amazon_datapulse/fraud_report/2024-01-15/theorchard/US' ), ) def test_fetch_from_athena_with_partition( self, mock_task_status, mock_set_overall_status, mock_athena, mock_s3_tasks, snapshot, ): """Test fetch_from_athena appends partition clauses and params.""" partition = {'cadence': 'daily', 'region': 'us-east-1'} tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=self.s3_path, partition=partition, ) call_kwargs = mock_athena.run_query.call_args.kwargs snapshot.assert_match(call_kwargs['athena_query'], 'athena_query.sql') assert call_kwargs['parameters'] == ( "'2024-01-15'", "'daily'", "'us-east-1'", ) def test_fetch_from_athena_cleans_s3_before_query( self, mock_task_status, mock_set_overall_status, mock_athena, mock_s3_tasks, ): """Test fetch_from_athena removes existing S3 files before query.""" tasks.fetch_from_athena( MagicMock(), date=_date, feed_name=self.feed_name, report_name=self.report_name, destination_s3_bucket=self.bucket, destination_s3_path=self.s3_path, ) mock_s3_tasks.remove_files_from_path.assert_called_once() remove_kwargs = mock_s3_tasks.remove_files_from_path.call_args.kwargs assert self.bucket in remove_kwargs['path'] assert self.s3_path in remove_kwargs['path']