"""Garcon tasks tests."""

import datetime
from unittest.mock import MagicMock
from unittest.mock import Mock
from unittest.mock import patch

from garcon_contrib.dynamo_feed_status import garcon_feed_status
import pytest

from flows import config
from flows import tasks


TEST_DATE = '2017-10-01'
YESTERDAY_DATE = (
    datetime.date.today() - datetime.timedelta(days=1)).strftime('%Y-%m-%d')


@pytest.yield_fixture
def mock_executor_context():
    """Yield executor context."""
    with patch('flows.tasks.SnowflakeSQLExecutorMR') as sf_executor:
        executor_context = sf_executor.return_value.__enter__.return_value
        yield executor_context


def test_bootstrap_task():
    """Test bootstrap task success."""
    result = tasks.bootstrap(Mock(), TEST_DATE)
    assert result['date'] == TEST_DATE


def test_bootstrap_task_no_date_provided():
    """Test bootstrap task when no date provided."""
    result = tasks.bootstrap(Mock(), None)
    assert result['date'] == YESTERDAY_DATE


def test_drop_create_tmp_table(mock_executor_context):
    """Test drop and create temporary table success."""
    mock_executor_context.drop_tmp_table.return_value = MagicMock()
    mock_executor_context.create_tmp_table.return_value = MagicMock()
    tasks.drop_create_temp_table(Mock())
    mock_executor_context.drop_tmp_table.assert_called()
    mock_executor_context.create_tmp_table.assert_called()


def test_drop_create_tmp_table_failure(mock_executor_context):
    """Test drop and create temporary table failure."""
    mock_executor_context.drop_tmp_table.side_effect = Exception('DbError')
    mock_executor_context.create_tmp_table.return_value = MagicMock()
    with pytest.raises(Exception, message='DbError'):
        tasks.drop_create_temp_table(Mock())
    mock_executor_context.create_tmp_table.assert_not_called()


def test_import_from_s3(mock_executor_context):
    """Test create stage and import from S3."""
    s3_path = 's3://test_bucket/{}'.format(TEST_DATE)
    mock_executor_context.create_import_stage.return_value = MagicMock()
    mock_executor_context.import_from_s3.return_value = MagicMock()
    tasks.import_from_s3(Mock(), TEST_DATE)
    mock_executor_context.create_import_stage.assert_called_with(
        s3_path, config.SF_IMPORT_STAGE_NAME)
    mock_executor_context.import_from_s3.assert_called()


def test_import_from_s3_create_stage_failed(mock_executor_context):
    """Test create stage failed during S3 import."""
    msg = 'DbError'
    mock_executor_context.create_import_stage.side_effect = Exception(msg)
    mock_executor_context.import_from_s3.return_value = MagicMock()
    with pytest.raises(Exception, message=msg):
        tasks.import_from_s3(Mock(), TEST_DATE)
    mock_executor_context.import_from_s3.assert_not_called()


def test_import_from_s3_failed(mock_executor_context):
    """Test S3 import failed."""
    msg = 'DbError'
    s3_path = 's3://test_bucket/{}'.format(TEST_DATE)
    mock_executor_context.create_import_stage.return_value = MagicMock()
    mock_executor_context.import_from_s3.side_effect = Exception(msg)
    with pytest.raises(Exception, message=msg):
        tasks.import_from_s3(Mock(), TEST_DATE)
    mock_executor_context.create_import_stage.assert_called_with(
        s3_path, config.SF_IMPORT_STAGE_NAME)


def test_merge_into_mr_table(mock_executor_context):
    """Test merge into MR table."""
    mock_executor_context.merge_into_mr_table.return_value = MagicMock()
    tasks.merge_into_mr_table(Mock())
    mock_executor_context.merge_into_mr_table.assert_called()


def test_merge_into_mr_table_failed(mock_executor_context):
    """Test merge into MR table failed."""
    error = 'DbError'
    mock_executor_context.merge_into_mr_table.side_effect = Exception(error)
    with pytest.raises(Exception, message=error):
        tasks.merge_into_mr_table()
    mock_executor_context.merge_into_mr_table.assert_called()


@patch('flows.tasks.garcon_feed_status.set_overall_status')
def test_update_dynamo_status(feed_status_mock):
    """Test update the final status in DynamoDB."""
    tasks.update_dynamo_status(Mock(), TEST_DATE)
    feed_status_mock.assert_called_with(
        config.SWF_NAME, TEST_DATE, garcon_feed_status.STATUS_INGESTED)


@patch('flows.tasks.garcon_feed_status.set_overall_status')
def test_update_dynamo_status_failed(feed_status_mock):
    """Test update the final status in DynamoDB failed."""
    msg = 'Amazon Error'
    feed_status_mock.side_effect = Exception(msg)
    with pytest.raises(Exception, message=msg):
        tasks.update_dynamo_status(Mock(), TEST_DATE)


def test_create_flattened_mr_table(mock_executor_context):
    """Test create flattened registry table."""
    mock_executor_context.create_flattened_mr_table.return_value = MagicMock()
    tasks.create_flattened_mr_table(Mock())
    mock_executor_context.create_flattened_mr_table.assert_called()


def test_create_flattened_mr_table_failed(mock_executor_context):
    """Test create flattened registry table failed."""
    error = 'DbError'
    mock_executor_context.create_flattened_mr_table.side_effect = Exception(
        error)
    with pytest.raises(Exception, message=error):
        tasks.create_flattened_mr_table()
    mock_executor_context.create_flattened_mr_table.assert_called()


def test_fill_flattened_mr_table(mock_executor_context):
    """Test fill flattened registry table."""
    mock_executor_context.fill_flattened_mr_table.return_value = MagicMock()
    mock_executor_context.truncate_table.return_value = MagicMock()
    tasks.fill_flattened_mr_table(Mock())
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.fill_flattened_mr_table.assert_called()


def test_fill_flattened_mr_table_failed(mock_executor_context):
    """Test fill flattened registry table failed."""
    error = 'DbError'
    mock_executor_context.truncate_table.return_value = MagicMock()
    mock_executor_context.fill_flattened_mr_table.side_effect = Exception(
        error)
    with pytest.raises(Exception, message=error):
        tasks.fill_flattened_mr_table()
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.create_flattened_mr_table.assert_called()


def test_fill_flattened_mr_locked_table(mock_executor_context):
    """Test fill flattened locked registry table."""
    mock_executor_context.fill_flattened_mr_locked_table.return_value = \
        MagicMock()
    mock_executor_context.truncate_table.return_value = MagicMock()
    tasks.fill_flattened_mr_locked_table(Mock())
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.fill_flattened_mr_locked_table.assert_called()


def test_fill_flattened_mr_locked_table_failed(mock_executor_context):
    """Test fill flattened locked registry table failed."""
    error = 'DbError'
    mock_executor_context.truncate_table.return_value = MagicMock()
    mock_executor_context.fill_flattened_mr_locked_table.side_effect = \
        Exception(error)
    with pytest.raises(Exception, message=error):
        tasks.fill_flattened_mr_locked_table()
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.create_flattened_mr_locked_table.assert_called()


def test_import_audit_from_s3(mock_executor_context):
    """Test create stage and import audit from S3."""
    s3_path = 's3://test_bucket/test_audit/{}'.format(TEST_DATE)
    mock_executor_context.create_import_stage.return_value = MagicMock()
    mock_executor_context.import_audit_from_s3.return_value = MagicMock()
    tasks.import_audit_from_s3(Mock(), TEST_DATE)
    mock_executor_context.create_import_stage.assert_called_with(
        s3_path, config.SF_AUDIT_IMPORT_STAGE_NAME)
    mock_executor_context.import_audit_from_s3.assert_called()


def test_import_audit_from_s3_create_stage_failed(mock_executor_context):
    """Test create stage failed during S3 import of MR_AUDIT."""
    msg = 'DbError'
    mock_executor_context.create_import_stage.side_effect = Exception(
        msg)
    mock_executor_context.import_audit_from_s3.return_value = MagicMock()
    with pytest.raises(Exception, message=msg):
        tasks.import_audit_from_s3(Mock(), TEST_DATE)
    mock_executor_context.import_audit_from_s3.assert_not_called()


def test_import_audit_from_s3_failed(mock_executor_context):
    """Test S3 import MR_AUDIT failed."""
    msg = 'DbError'
    s3_path = 's3://test_bucket/test_audit/{}'.format(TEST_DATE)
    mock_executor_context.create_import_stage.return_value = MagicMock()
    mock_executor_context.import_audit_from_s3.side_effect = Exception(msg)
    with pytest.raises(Exception, message=msg):
        tasks.import_audit_from_s3(Mock(), TEST_DATE)
    mock_executor_context.create_import_stage.assert_called_with(
        s3_path, config.SF_AUDIT_IMPORT_STAGE_NAME)


def test_fill_flattened_mr_audit_table(mock_executor_context):
    """Test fill flattened audit registry table."""
    mock_executor_context.fill_flattened_mr_audit_table.return_value = \
        MagicMock()
    mock_executor_context.truncate_table.return_value = MagicMock()
    tasks.fill_flattened_mr_audit_table(Mock())
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.fill_flattened_mr_audit_table.assert_called()


def test_fill_flattened_mr_laudit_table_failed(mock_executor_context):
    """Test fill flattened locked registry table failed."""
    error = 'DbError'
    mock_executor_context.truncate_table.return_value = MagicMock()
    mock_executor_context.fill_flattened_mr_audit_table.side_effect = \
        Exception(error)
    with pytest.raises(Exception, message=error):
        tasks.fill_flattened_mr_audit_table()
    mock_executor_context.truncate_table.assert_called()
    mock_executor_context.fill_flattened_mr_audit_table.assert_called()