"""Tests for MRSnowflake Executor."""

from unittest.mock import MagicMock
from unittest.mock import patch

import pytest
from snowflake import connector

from flows import config
from flows.snowflake_executor import SnowflakeSQLExecutorMR


@pytest.fixture
def mock_snowflake_connector(monkeypatch):
    """Mock low-level snowflake connector object."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    return connect_mock


def test_create_tmp_table(mock_snowflake_connector):
    """Test create temporary table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.create_tmp_table()
    test_table_name = '{db}.{schema}.{table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        table=config.SF_TEMP_TABLE_NAME
    )
    assert expatch.called
    assert test_table_name in expatch.call_args[0][0]


def test_create_tmp_table_failed(mock_snowflake_connector):
    """Test create temporary table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.create_tmp_table()
    assert expatch.called


def test_create_import_stage(mock_snowflake_connector):
    """Test create import stage."""
    executor = SnowflakeSQLExecutorMR()
    import_date = '2017-01-01'
    s3_path = 's3://{bucket_name}/{path}'.format(
        bucket_name=config.S3_BUCKET_PATH,
        path=import_date
    )
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.create_import_stage(s3_path, config.SF_IMPORT_STAGE_NAME)
    import_stage_name = '{db}.{schema}.{stage}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        stage=config.SF_IMPORT_STAGE_NAME
    )
    assert expatch.called
    assert import_stage_name in expatch.call_args[0][0]
    assert s3_path == expatch.call_args[0][1]['s3_location']


def test_create_import_stage_failed(mock_snowflake_connector):
    """Test create import stage failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.create_import_stage(
                's3://path', config.SF_IMPORT_STAGE_NAME)
    assert expatch.called


def test_import_from_s3(mock_snowflake_connector):
    """Test import from S3."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.import_from_s3()

    test_table_name = '{db}.{schema}.{table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        table=config.SF_TEMP_TABLE_NAME
    )

    import_stage_name = '{db}.{schema}.{stage}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        stage=config.SF_IMPORT_STAGE_NAME
    )

    assert expatch.called
    assert test_table_name in expatch.call_args[0][0]
    assert import_stage_name in expatch.call_args[0][0]


def test_import_from_s3_failed(mock_snowflake_connector):
    """Test import from S3 failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.import_from_s3()
    assert expatch.called


def test_merge_into_mr_table(mock_snowflake_connector):
    """Test merge into MR table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.merge_into_mr_table()

    tmp_table_name = '{db}.{schema}.{table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        table=config.SF_TEMP_TABLE_NAME
    )
    mr_table_name = '{db}.{schema}.{mr_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_table=config.SF_TEMP_TABLE_NAME
    )
    grouped_view_name = config.SF_GROUPED_VIEW_NAME
    sql = expatch.call_args[0][0]

    assert expatch.called
    assert tmp_table_name in sql
    assert mr_table_name in sql
    assert grouped_view_name in sql


def test_merge_into_mr_table_failed(mock_snowflake_connector):
    """Test merge into MR table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.merge_into_mr_table()
    assert expatch.called


def test_create_flattened_mr_table(mock_snowflake_connector):
    """Test create flattened registry table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.create_flattened_mr_table()

    mr_flattened_table_name = '{db}.{schema}.{mr_flattened_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_flattened_table=config.SF_MR_FLATTENED_TABLE_NAME
    )

    assert expatch.called
    sql = expatch.call_args[0][0]
    assert mr_flattened_table_name in sql


def test_create_flattened_mr_table_failed(mock_snowflake_connector):
    """Test create flattened registry table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.create_flattened_mr_table()
    assert expatch.called


def test_fill_flattened_mr_table(mock_snowflake_connector):
    """Test fill flattened refistry table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.fill_flattened_mr_table()

    mr_flattened_table_name = '{db}.{schema}.{mr_flattened_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_flattened_table=config.SF_MR_FLATTENED_TABLE_NAME
    )

    mr_table_name = '{db}.{schema}.{mr_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_table=config.SF_MR_TABLE_NAME
    )

    assert expatch.called
    sql = expatch.call_args[0][0]
    assert mr_flattened_table_name in sql
    assert mr_table_name in sql


def test_fill_flattened_mr_table_failed(mock_snowflake_connector):
    """Test fill flattened registry table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.fill_flattened_mr_table()
    assert expatch.called


def test_fill_flattened_mr_locked_table(mock_snowflake_connector):
    """Test fill flattened locked territories registry table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.fill_flattened_mr_locked_table()

    mr_locked_flattened_table_name = '{db}.{schema}.{mr_locked_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_locked_table=config.SF_MR_LOCKED_FLATTENED_TABLE_NAME
    )

    mr_table_name = '{db}.{schema}.{mr_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_table=config.SF_MR_TABLE_NAME
    )

    assert expatch.called
    sql = expatch.call_args[0][0]
    assert mr_locked_flattened_table_name in sql
    assert mr_table_name in sql


def test_fill_flattened_mr_locked_table_failed(mock_snowflake_connector):
    """Test fill flattened locked territories registry table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.fill_flattened_mr_locked_table()
    assert expatch.called


def test_import_audit_from_s3(mock_snowflake_connector):
    """Test import audit table from S3."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.import_audit_from_s3()

    test_table_name = '{db}.{schema}.{table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        table=config.SF_MR_AUDIT_TABLE_NAME
    )

    import_stage_name = '{db}.{schema}.{stage}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        stage=config.SF_AUDIT_IMPORT_STAGE_NAME
    )

    assert expatch.called
    assert test_table_name in expatch.call_args[0][0]
    assert import_stage_name in expatch.call_args[0][0]


def test_import_audit_from_s3_failed(mock_snowflake_connector):
    """Test import audit table from S3 failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.import_audit_from_s3()
    assert expatch.called


def test_fill_flattened_mr_audit_table(mock_snowflake_connector):
    """Test fill flattened audit registry table."""
    executor = SnowflakeSQLExecutorMR()
    with patch.object(executor, 'execute', wraps=executor.execute) as expatch:
        executor.fill_flattened_mr_audit_table()

    mr_audit_flattened_table_name = '{db}.{schema}.{mr_audit_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_audit_table=config.SF_MR_AUDIT_FLATTENED_TABLE_NAME
    )

    mr_audit_table_name = '{db}.{schema}.{mr_audit_table}'.format(
        db=config.SF_CONFIG['db'],
        schema=config.SF_CONFIG['schema'],
        mr_audit_table=config.SF_MR_AUDIT_TABLE_NAME
    )

    assert expatch.called
    sql = expatch.call_args[0][0]
    assert mr_audit_flattened_table_name in sql
    assert mr_audit_table_name in sql


def test_fill_flattened_mr_audit_table_failed(mock_snowflake_connector):
    """Test fill flattened audit registry table failed."""
    executor = SnowflakeSQLExecutorMR()
    error_msg = 'DbError'
    with patch.object(
            executor, 'execute', side_effect=Exception(error_msg)) as expatch:
        with pytest.raises(Exception, msg=error_msg):
            executor.fill_flattened_mr_audit_table()
    assert expatch.called