from unittest.mock import MagicMock

from feed_ingestion.common.base_executor import SnowflakeTemplatedSQLExecutor
from feed_ingestion.flows.amazon_datapulse import flow
from feed_ingestion.flows.amazon_datapulse import stage_loader


def test_delete_from_staging_raw_template(snapshot, monkeypatch):
    monkeypatch.setattr(
        SnowflakeTemplatedSQLExecutor, 'get_connection', MagicMock()
    )
    sf_config = {
        'db': 'TEST_DB',
        'schema': 'TEST_SCHEMA',
    }
    executor = SnowflakeTemplatedSQLExecutor(sf_config=sf_config)
    stage_loader_obj = stage_loader.AmazonDataPulseSL(
        executor=executor,
        sql_loader=flow.sql_loader,
    )
    execute_mock = MagicMock()
    monkeypatch.setattr(executor, 'execute', execute_mock)

    stage_loader_obj.clean_staging_raw_table(
        staging_raw_table='sample',
        date='2024-01-01',
        partition={
            'report_date': '2025-01-01',
            'some_int': 13,
        },
        staging_raw_database='TEST_DB',
        staging_raw_schema='TEST_SCHEMA',
    )

    execute_mock.assert_called()
    execute_args = execute_mock.call_args[1]

    snapshot.assert_match(
        execute_args['sql_template'], 'delete_from_staging_raw.sql'
    )
    assert execute_args['params'] == {
        'date_1': '2024-01-01',
        'val_2': '2025-01-01',
        'val_3': 13,
    }


def test_load_staging_raw_table(snapshot, monkeypatch):
    monkeypatch.setattr(
        SnowflakeTemplatedSQLExecutor, 'get_connection', MagicMock()
    )
    sf_config = {
        'db': 'TEST_DB',
        'schema': 'TEST_SCHEMA',
    }
    executor = SnowflakeTemplatedSQLExecutor(sf_config=sf_config)
    stage_loader_obj = stage_loader.AmazonDataPulseSL(
        executor=executor,
        sql_loader=flow.sql_loader,
    )
    execute_mock = MagicMock()
    monkeypatch.setattr(executor, 'execute', execute_mock)

    stage_loader_obj.load_staging_raw_table(
        staging_raw_table='staging_raw_fraud_report',
        source_files_dict={},
        date='2024-01-15',
        stage_name='test_stage',
        query_name='load_staging_raw_fraud_report',
        staging_raw_database='TEST_DB',
        staging_raw_schema='TEST_SCHEMA',
    )

    execute_mock.assert_called_once()
    execute_args = execute_mock.call_args[1]

    snapshot.assert_match(
        execute_args['sql_template'], 'load_staging_raw_fraud_report.sql'
    )
    assert execute_args['params'] == {
        'download_date_1': '2024-01-15',
        'pattern_2': r'.*/\d{8}_\d{6}_\d{5}_\w+_\w+\-\w+\-\w+\-\w+\-\w+',
    }