"""Unit tests for tasks for the Chartmetric Tracks Ingestion Workflow.""" import datetime from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import patch from freezegun import freeze_time import pytest from feed_ingestion.flows.chartmetric_tracks import config from feed_ingestion.flows.chartmetric_tracks import tasks @pytest.fixture def feed_name(): """Feed name value.""" return 'chartmetric_tracks' @pytest.fixture def neo4j_config(mocker): """Mock get_neo4j_config config.""" return mocker.patch.object(tasks, 'get_neo4j_config') @pytest.fixture def neo4j_executor_class(mocker): """Mock Neo4jExecutor class.""" return mocker.patch.object(tasks, 'Neo4jExecutor') @pytest.fixture def neo4j_executor(neo4j_executor_class): """Mock Neo4jExecutor context manger instance.""" with neo4j_executor_class.return_value as neo4j_executor: yield neo4j_executor @freeze_time('2020-01-01') def test_bootstrap_default_date(monkeypatch, feed_name): """Test that bootstrap returns the default date.""" date = None date_limit = None expected = { 'feed_name': feed_name, 'date': '2019-12-31', 'date_limit': '3000-01-01', 'full_refresh': False, 'platform_names': config.all_platform_names } result = tasks.bootstrap(MagicMock(), date, date_limit) assert result == expected @freeze_time('2020-01-01') def test_bootstrap_context_date(monkeypatch, feed_name): """Test that bootstrap returns the context date.""" date = '2019-01-01' date_limit = '2019-01-02' full_refresh = 'True' expected = { 'feed_name': feed_name, 'date': date, 'date_limit': date_limit, 'full_refresh': True, 'platform_names': config.all_platform_names } result = tasks.bootstrap( MagicMock(), date, full_refresh, date_limit) assert result == expected @patch('feed_ingestion.flows.chartmetric_tracks.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.SnowflakeExecutor') def test_create_table(mock_executor_class, mock_get_sf_config, feed_name): """Test create table.""" mock_conifg = {} mock_get_sf_config.return_value = mock_conifg mock_executor = MagicMock() mock_executor_class.return_value.__enter__.return_value = mock_executor date = '2020-01-01' date_limit = '3000-01-01' platform_name = 'spotify' full_refresh = False tasks.create_table( MagicMock(), feed_name, date, date_limit, full_refresh, platform_name) mock_get_sf_config.assert_called_once_with(feed_name) mock_executor_class.assert_called_once_with(mock_conifg) mock_executor.execute_query.assert_called_once_with( 'create_spotify_table', table_name='chartmetric_tracks_spotify', ingest_date=date, ingest_date_limit=date_limit, platform_name=platform_name, full_refresh=full_refresh) @patch('feed_ingestion.flows.chartmetric_tracks.tasks._ingest_data') def test_ingest_data_task(_ingest_data_mock, feed_name): """Test ingest data task.""" platform_names = ['amazon', 'spotify'] activity_mock = MagicMock() tasks.ingest_data(activity_mock, feed_name, platform_names) _ingest_data_mock.assert_has_calls([ call(activity_mock, feed_name, 'amazon'), call(activity_mock, feed_name, 'spotify') ]) @patch('feed_ingestion.flows.chartmetric_tracks.tasks.SnowflakeExecutor') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.get_neo4j_config') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.Neo4jExecutor') def test_ingest_data( mock_neo4j_executor_class, mock_get_sf_config, mock_get_neo4j_config, mock_sf_executor_class, feed_name): """Test ingest data logic.""" mock_neo4j_config = {} mock_get_neo4j_config.return_value = mock_neo4j_config mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_neo4j_executor = MagicMock() mock_neo4j_executor_class.return_value.__enter__.return_value = \ mock_neo4j_executor mock_sf_executor = MagicMock() mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor mock_sf_executor.fetchall_dict_query.side_effect = [[{}], []] platform_name = 'spotify' tasks._ingest_data(MagicMock(), feed_name, platform_name) call_list = [ call( 'get_spotify_raw_data', limit=config.row_limit_per_batch, offset=0, table_name='chartmetric_tracks_spotify'), call( 'get_spotify_raw_data', limit=config.row_limit_per_batch, offset=config.row_limit_per_batch, table_name='chartmetric_tracks_spotify') ] mock_get_sf_config.assert_called_once_with(feed_name) mock_sf_executor.fetchall_dict_query.assert_has_calls(call_list) mock_get_neo4j_config.assert_called_once_with(feed_name) mock_neo4j_executor_class.assert_called_once_with( feed_name, mock_neo4j_config) mock_neo4j_executor.execute_write_query.assert_called_once_with( 'ingest_spotify_raw_data', {'rows': [{}]}) @patch('feed_ingestion.flows.chartmetric_tracks.tasks.time') def test_write_delay_first_call(time_mock): """Test write delay first call.""" tasks._write_delay(MagicMock(), None) time_mock.sleep.assert_not_called() @freeze_time('2024-01-01 00:00:30') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.time') def test_write_delay_paused(time_mock): """Test write delay.""" last_write_at = ( datetime.datetime.now() - datetime.timedelta(seconds=config.WRITE_TRANSACTION_INTERVAL - 2)) tasks._write_delay(MagicMock(), last_write_at) assert time_mock.sleep.called @freeze_time('2024-01-01 00:00:30') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.time') def test_write_delay_not_paused(time_mock): """Test write delay.""" last_write_at = ( datetime.datetime.now() - datetime.timedelta(seconds=config.WRITE_TRANSACTION_INTERVAL + 2)) tasks._write_delay(MagicMock(), last_write_at) assert not time_mock.sleep.called @patch('feed_ingestion.flows.chartmetric_tracks.tasks.get_neo4j_config') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.get_sf_config') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.Neo4jExecutor') @patch('feed_ingestion.flows.chartmetric_tracks.tasks.SnowflakeExecutor') def test_delete_defunct_participations( mock_sf_executor_class, mock_neo4j_executor_class, mock_get_sf_config, mock_get_neo4j_config): """Test delete_defunct_participations task.""" mock_get_neo4j_config.return_value = {} mock_sf_config = {'db': 'DB', 'schema': 'SCHEMA'} mock_get_sf_config.return_value = mock_sf_config mock_neo4j_executor = MagicMock() mock_neo4j_executor_class.return_value.__enter__.return_value = \ mock_neo4j_executor mock_sf_executor = MagicMock() first_batch = [{'key': 'value'}] second_batch = [] mock_sf_executor.fetchall_dict_query.side_effect = \ (first_batch, second_batch) mock_sf_executor_class.return_value.__enter__.return_value = \ mock_sf_executor sf_params_expected_1 = { 'limit': config.row_limit_per_batch, 'offset': 0, } sf_params_expected_2 = { 'limit': config.row_limit_per_batch, 'offset': config.row_limit_per_batch, } tasks.delete_defunct_participations( MagicMock(), 'some_feed_name', '2020-01-01') sf_query_name = 'delete_defunct_participations' neo4j_query_name = 'delete_defunct_participations' assert mock_sf_executor.fetchall_dict_query.call_args_list == [ call(sf_query_name, **sf_params_expected_1), call(sf_query_name, **sf_params_expected_2), ] sf_params_expected_1 = {'rows': first_batch} assert mock_neo4j_executor.execute_write_query.call_args_list == [ call(neo4j_query_name, sf_params_expected_1), ]