"""Unit tests for the generators and util functions of the sql2sf workflow."""

from unittest.mock import MagicMock

from snowflake_etl.flows.sql2sf import config
from snowflake_etl.flows.sql2sf import generators
from snowflake_etl.flows.sql2sf import helpers


def test__get_min_max_pk(monkeypatch):
    """Test _get_min_max_pk helper."""
    execute_mock = MagicMock(return_value=(1, 100))
    monkeypatch.setattr(helpers, 'execute_with_mysql', execute_mock)
    min_id, max_id = generators._get_min_max_pk(
        'some_key', 'some_schema', 'some_table', ('', '', '', ''))
    assert min_id == 1
    assert max_id == 100
    execute_mock.assert_called_with(
        '\n        SELECT MIN(some_key) as min_id,\n               '
        'MAX(some_key) as max_id\n        FROM `some_schema`.`some_table`\n  '
        '      ', 'fetchone', '', '', '', '')


def test_get_min_max_pairs(monkeypatch):
    """Test get_min_max_pairs helper."""
    _get_min_max_mock = MagicMock(return_value=(1, 100))
    monkeypatch.setattr(generators, '_get_min_max_pk', _get_min_max_mock)
    result = generators.get_min_max_pairs(
        'some_key', 50, 'some_schema', 'some_table',  ('', '', '', ''))

    assert result == [
        {'max_id': 50, 'min_id': 1}, {'max_id': 100, 'min_id': 51}]


def test_get_column_names(monkeypatch):
    """Test get_columns helper."""
    execute_mock = MagicMock(
        return_value=[('somecolumn1', 'varchar'), ('somecolumn2', 'int')])
    monkeypatch.setattr(helpers, 'execute_with_mysql', execute_mock)
    result = generators.get_columns(
        'art_relations', 'vendor', ('', '', '', ''))
    assert result == [('somecolumn1', True), ('somecolumn2', None)]
    execute_mock.assert_called_with(
        'SELECT COLUMN_NAME, DATA_TYPE\n        '
        'FROM information_schema.columns\n        '
        "WHERE TABLE_SCHEMA = 'art_relations'\n          "
        "AND TABLE_NAME = 'vendor'\n        "
        'ORDER BY ORDINAL_POSITION;'
        '', 'fetchall', '', '', '', '')


def test_unload_mysql_data_params_not_in_chunks(monkeypatch):
    """Test unload_mysql_data_params generator (no chunks)."""
    context = {
        'bootstrap_sql2sf.source_db_host': 'some_host',
        'bootstrap_sql2sf.source_schema': 'some_schema',
        'bootstrap_sql2sf.source_table': 'some_table',
        'bootstrap_sql2sf.destination_s3_bucket': 'some_bucket',
        'bootstrap_sql2sf.destination_s3_key': 'some_key'
    }

    execute_mock = MagicMock(
        return_value=[('somecolumn1', 'varchar'), ('somecolumn2', 'int')])
    monkeypatch.setattr(helpers, 'execute_with_mysql', execute_mock)
    sources_mock_obj = MagicMock()
    sources_mock_obj.is_unload_in_chunks = MagicMock(return_value=False)
    sources_mock_obj.get_db_credentials = MagicMock(side_effect=[()])
    sources_mock = MagicMock(return_value=sources_mock_obj)
    monkeypatch.setattr(config, 'SourcesConf', sources_mock)

    result = list(generators.unload_mysql_data_params(context))[0]
    assert result == {
        'generator.bucket': 'some_bucket',
        'generator.destination_s3_key': 'some_key',
        'generator.columns': [('somecolumn1', True), ('somecolumn2', None)],
        'generator.table': 'some_table'}


def test_unload_mysql_data_params_in_chunks(monkeypatch):
    """Test unload_mysql_data_params generator (chunked unload)."""
    context = {
        'bootstrap_sql2sf.source_db_host': 'some_host',
        'bootstrap_sql2sf.source_schema': 'some_schema',
        'bootstrap_sql2sf.source_table': 'some_table',
        'bootstrap_sql2sf.destination_s3_bucket': 'some_bucket',
        'bootstrap_sql2sf.destination_s3_key': 'some_key'
    }

    execute_mock = MagicMock(
        return_value=[('somecolumn1', 'varchar'), ('somecolumn2', 'int')])
    monkeypatch.setattr(helpers, 'execute_with_mysql', execute_mock)
    sources_mock_obj = MagicMock()
    sources_mock_obj.is_unload_in_chunks = MagicMock(return_value=True)
    sources_mock_obj.get_db_credentials = MagicMock(side_effect=[()])
    sources_mock_obj.get_chunk_size = MagicMock(return_value=50)
    sources_mock_obj.get_primary_key = MagicMock(return_value='some_key')
    sources_mock = MagicMock(return_value=sources_mock_obj)
    monkeypatch.setattr(config, 'SourcesConf', sources_mock)

    get_min_max_list_mock = MagicMock(return_value=[
        {'max_id': 50, 'min_id': 1}, {'max_id': 100, 'min_id': 51}])
    monkeypatch.setattr(
        generators, 'get_min_max_pairs', get_min_max_list_mock)

    results = list(generators.unload_mysql_data_params(context))
    assert results[0] == {
        'generator.bucket': 'some_bucket',
        'generator.destination_s3_key': 'some_key',
        'generator.table': 'some_table',
        'generator.columns': [('somecolumn1', True), ('somecolumn2', None)],
        'generator.min_id': 1,
        'generator.max_id': 50}
    assert results[1] == {
        'generator.bucket': 'some_bucket',
        'generator.destination_s3_key': 'some_key',
        'generator.table': 'some_table',
        'generator.columns': [('somecolumn1', True), ('somecolumn2', None)],
        'generator.min_id': 51,
        'generator.max_id': 100}