"""Unit tests for generic tasks of the FlowLoadRawMixinSF mixin."""

from unittest.mock import call, MagicMock, patch

from snowflake import connector

from feed_ingestion.tasks import load_raw_table_tasks_sf
from feed_ingestion.util import task_status


def test_create_temp_staging_raw_table(monkeypatch, sf_config_mock):
    """Test create_temp_staging_raw_table task."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    monkeypatch.setattr(task_status, 'get_report_contexts', MagicMock(
        return_value={}))
    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.create_temp_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed', 'temp_table',
            sf_config_mock)

    executor_with_mock.drop_table.assert_called_with('temp_table')
    executor_with_mock.create_temp_staging_raw_table. \
        assert_called_with('temp_table')


def test_load_temp_staging_raw_table(monkeypatch, sf_config_mock):
    """Test load_temp_staging_raw_table task."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    aws = {'access_key': 'test_access', 'access_secret': 'test_secret'}
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.load_temp_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed',
            aws, 's3keydir/path', 'temp_table', sf_config_mock)

    executor_with_mock.load_temp_staging_raw_table.assert_called_with(
        'temp_table', aws, 's3keydir/path')


def test_load_staging_raw_table_with_clean(monkeypatch, sf_config_mock):
    """Test load_staging_raw_table task."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    set_complete = MagicMock()
    monkeypatch.setattr(task_status, 'mark_completed_task', set_complete)

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.load_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed', 'temp_table',
            'perm_table', sf_config_mock, clean='True')

    executor_with_mock.clean_staging_raw_table.assert_called_with(
        'perm_table', '2017-08-02')
    executor_with_mock.load_staging_raw_table.assert_called_with(
        'temp_table', 'perm_table', '2017-08-02')
    executor_with_mock.drop_table.assert_called_with('temp_table')


def test_drop_temp_staging_raw_table(monkeypatch, sf_config_mock):
    """Test drop_temp_staging_raw_table task."""
    activity_mock = MagicMock()
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.drop_temp_staging_raw_table(
            activity_mock, 'test_feed', 'temp_table', sf_config_mock)

    executor_with_mock.drop_table.assert_called_with('temp_table')


def test_drop_remaining_temp_staging_raw_tables(monkeypatch, sf_config_mock):
    """Test load_staging_raw_table task set complete."""
    activity_mock = MagicMock()
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    metadata = ['metadata' for _ in range(13)]

    executor_with_mock = MagicMock()
    executor_with_mock.select_remaining_tables.return_value = [
        ('created_on', 'temp_staging_pattern_20170802', 'TEST', 'SCHEMA',
         *metadata),
        ('created_on', 'temp_staging_pattern_meta_20170802', 'TEST', 'SCHEMA',
         *metadata),
        ('created_on', 'temp_staging_pattern_meta2_20170802', 'TEST', 'SCHEMA',
         *metadata),
        ('created_on', 'staging_raw_pattern_20170802', 'TEST', 'SCHEMA',
         *metadata)  # hypothetical Snowflake SHOW TABLES output issue
    ]
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock

    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.drop_remaining_temp_staging_raw_tables(
            activity_mock, 'test_feed', 'temp_staging_pattern',
            '2017-08-02', sf_config_mock)

    for t in ['temp_staging_pattern_meta2_20170802',
              'temp_staging_pattern_meta_20170802',
              'temp_staging_pattern_20170802', ]:
        assert call(t) in executor_with_mock.drop_table.mock_calls

    assert call('staging_raw_pattern_20170802') not in \
           executor_with_mock.drop_table.mock_calls


def test_skip_cleanup_staging_raw_table(monkeypatch, sf_config_mock):
    """Test load_staging_raw_table task skips staging_raw cleanup."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.load_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed', 'temp_table',
            'perm_table', sf_config_mock, clean='False')

    executor_with_mock.load_staging_raw_table.assert_called_with(
        'temp_table', 'perm_table', '2017-08-02')
    executor_with_mock.drop_table.assert_called_with('temp_table')


def test_load_staging_raw_table_set_complete(monkeypatch, sf_config_mock):
    """Test load_staging_raw_table task set complete."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    set_complete = MagicMock()
    monkeypatch.setattr(task_status, 'mark_completed_task', set_complete)

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.load_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed', 'temp_table',
            'perm_table', sf_config_mock, set_complete='True')

    set_complete.assert_called_with(
        'test_feed', '2017-08-02', 'staging_raw_table_tasks')


def test_load_staging_raw_table_skip_set_complete(monkeypatch, sf_config_mock):
    """Test load_staging_raw_table task set complete."""
    activity_mock = MagicMock()
    monkeypatch.setattr(task_status, 'is_completed_task', MagicMock(
        return_value=False))
    monkeypatch.setattr(task_status, 'mark_completed_task', MagicMock())
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    set_complete = MagicMock()
    monkeypatch.setattr(task_status, 'mark_completed_task', set_complete)

    executor_with_mock = MagicMock()
    executor_mock = MagicMock()
    executor_mock.__enter__.return_value = executor_with_mock
    executor_class_mock = MagicMock(return_value=executor_mock)
    with patch.dict(load_raw_table_tasks_sf.registered_executors,
                    {'test_feed': executor_class_mock}):
        load_raw_table_tasks_sf.load_staging_raw_table(
            activity_mock, '2017-08-02', 'test_feed', 'temp_table',
            'perm_table', sf_config_mock, set_complete='False')

    set_complete.assert_not_called()


def test_mark_staging_raw_table_tasks_complete(monkeypatch):
    """Test mark_staging_raw_table_tasks_complete task set complete."""
    activity_mock = MagicMock()

    set_complete = MagicMock()
    monkeypatch.setattr(task_status, 'mark_completed_task', set_complete)

    load_raw_table_tasks_sf.mark_staging_raw_table_tasks_complete(
        activity_mock, '2017-08-02', 'test_feed')

    set_complete.assert_called_with(
        'test_feed', '2017-08-02', 'staging_raw_table_tasks')