"""Unit tests for Amazon Music Ingestion Workflow."""
from unittest import mock
from unittest.mock import MagicMock

from feed_ingestion.flows.amazon_music.flow import Flow


def test_decider(monkeypatch):
    """Test normal decider execution."""
    monkeypatch.setenv('FEED_INGESTION_SNS_TOPIC', 'test_topic')
    schedule = MagicMock()
    schedule_result_object = MagicMock()
    schedule_result_object.result = {
        'update_dim_tables.sns_report_subject': 'Test subj'
    }
    schedule.return_value = schedule_result_object

    flow = Flow()
    flow.decider(schedule)

    schedule.assert_has_calls([
        mock.call('bootstrap', mock.ANY),
        mock.call('remove_stale_stage_files', mock.ANY, requires=[mock.ANY]),
        mock.call('grab_drop_files', mock.ANY, requires=[mock.ANY]),
        mock.call('update_feed_s3_file_status', mock.ANY, requires=[mock.ANY]),
        mock.call('set_feed_status_downloaded', mock.ANY, requires=[mock.ANY]),
        mock.call('clean_files', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'create_temp_staging_raw_tables', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'load_temp_staging_raw_tables', mock.ANY, requires=[mock.ANY]),
        mock.call('clean_staging_raw_table', mock.ANY, requires=mock.ANY),
        mock.call('load_staging_raw_table', mock.ANY, requires=mock.ANY),
        mock.call(
            'mark_staging_raw_table_tasks_complete',
            mock.ANY, requires=mock.ANY),
        mock.call(
            'set_status_populated_raw_table', mock.ANY, requires=mock.ANY),
        mock.call(
            'drop_temp_staging_raw_tables', mock.ANY, requires=mock.ANY),
        mock.call('load_aggregated_table', mock.ANY, requires=mock.ANY),
        mock.call('update_dim_tables', mock.ANY, requires=mock.ANY),
        mock.call(
            'send_dimension_tables_update_report_sns_notification', mock.ANY,
            requires=mock.ANY),
        mock.call('load_staging_fact_table', mock.ANY, requires=mock.ANY),
        mock.call(
            'load_aggregated_skips_and_saves', mock.ANY, requires=mock.ANY),
        mock.call('load_fact_tables', mock.ANY, requires=mock.ANY),
        mock.call('clean_up', mock.ANY, requires=mock.ANY),
        mock.call('build_jenkins_dbt', mock.ANY, requires=mock.ANY),
        mock.call('monitor_drop_location', mock.ANY, requires=mock.ANY)
    ])


def test_decider_should_stop_when_data_already_ingested():
    """Decider quits if data is already ingested."""
    bootstrap = MagicMock()
    bootstrap.result = {'bootstrap_feed.stop': True}
    schedule = MagicMock(return_value=bootstrap)

    flow = Flow()
    result = flow.decider(schedule)

    assert result is None
    schedule.assert_has_calls([
        mock.call('bootstrap', mock.ANY),
        mock.call('remove_stale_stage_files', mock.ANY, requires=[mock.ANY]),
    ])
    assert schedule.call_count == 2


def test_decider_datapulse(monkeypatch):
    """source=datapulse skips download chain and loads from Snowflake."""
    monkeypatch.setenv('FEED_INGESTION_SNS_TOPIC', 'test_topic')
    schedule_result_object = MagicMock()
    schedule_result_object.result = {
        'bootstrap_feed.source': 'datapulse',
        'update_dim_tables.sns_report_subject': 'Test subj',
    }
    schedule = MagicMock(return_value=schedule_result_object)

    flow = Flow()
    flow.decider(schedule)

    schedule.assert_has_calls([
        mock.call('bootstrap', mock.ANY),
        mock.call('clean_staging_raw_table', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'load_staging_raw_datapulse', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'mark_staging_raw_table_tasks_complete', mock.ANY,
            requires=[mock.ANY]),
        mock.call(
            'set_status_populated_raw_table', mock.ANY, requires=[mock.ANY]),
        mock.call('load_aggregated_table', mock.ANY, requires=[mock.ANY]),
        mock.call('update_dim_tables', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'send_dimension_tables_update_report_sns_notification', mock.ANY,
            requires=[mock.ANY]),
        mock.call(
            'load_staging_fact_table', mock.ANY,
            requires=[mock.ANY, mock.ANY]),
        mock.call(
            'load_aggregated_skips_and_saves', mock.ANY, requires=[mock.ANY]),
        mock.call(
            'load_fact_tables', mock.ANY, requires=[mock.ANY, mock.ANY]),
        mock.call('clean_up', mock.ANY, requires=[mock.ANY]),
        mock.call('build_jenkins_dbt', mock.ANY, requires=[mock.ANY]),
    ])
    # sftp-only steps must not be scheduled in datapulse mode
    scheduled_names = {call.args[0] for call in schedule.call_args_list}
    forbidden = {
        'remove_stale_stage_files',
        'grab_drop_files',
        'awal_grab_drop_files',
        'awal_grab_drop_files_s3',
        'altafonte_grab_drop_files',
        'altafonte_grab_drop_files_s3',
        'sme_grab_and_clean_files',
        'map_expected_files_to_sme',
        'update_feed_s3_file_status',
        'set_feed_status_downloaded',
        'convert_zip_to_gzip_on_s3',
        'clean_files',
        'create_temp_staging_raw_tables',
        'load_temp_staging_raw_tables',
        'load_staging_raw_table',
        'drop_temp_staging_raw_tables',
        'monitor_drop_location',
    }
    assert forbidden.isdisjoint(scheduled_names), (
        f'unexpected sftp steps in datapulse DAG: '
        f'{forbidden & scheduled_names}'
    )


def test_decider_datapulse_stop_short_circuits():
    """Datapulse branch respects bootstrap_feed.stop."""
    bootstrap = MagicMock()
    bootstrap.result = {
        'bootstrap_feed.source': 'datapulse',
        'bootstrap_feed.stop': True,
    }
    schedule = MagicMock(return_value=bootstrap)

    flow = Flow()
    result = flow.decider(schedule)

    assert result is None
    schedule.assert_called_once_with('bootstrap', mock.ANY)


def test_decider_datapulse_populate_only_stops_after_raw():
    """populate_only stops after set_status_populated_raw_table."""
    schedule_result_object = MagicMock()
    schedule_result_object.result = {
        'bootstrap_feed.source': 'datapulse',
        'bootstrap_feed.populate_only': True,
    }
    schedule = MagicMock(return_value=schedule_result_object)

    flow = Flow()
    flow.decider(schedule)

    scheduled_names = [call.args[0] for call in schedule.call_args_list]
    assert 'set_status_populated_raw_table' in scheduled_names
    assert 'load_aggregated_table' not in scheduled_names
    assert 'load_fact_tables' not in scheduled_names