"""Unit tests for Snowflake DB utility functions."""

from unittest.mock import MagicMock
from unittest.mock import call

from snowflake import connector

from analytics_aggregation.util.snowflake_db import SnowflakeSQLExecutor


def test_init(monkeypatch, sf_config_mock):
    """Test __init__ method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    SnowflakeSQLExecutor(sf_config=sf_config_mock)
    connect_mock.assert_called_with(
        warehouse='test_wh', password='test_pass', account='test_acc',
        schema='test_schema', user='test_user',
        db='test_db', role='test_role', autocommit=True)


def test_enter(monkeypatch, sf_config_mock):
    """Test __enter__ method to check if context manager implemented."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        assert isinstance(sf_exec, SnowflakeSQLExecutor)


def test_exit(monkeypatch, sf_config_mock):
    """Test __exit__ method to check if context manager implemented."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock):
        pass
    connect_mock.assert_has_calls([call().close()])


def test_get_connection(monkeypatch, sf_config_mock):
    """Test get_connection method returns the result of call of connect()."""
    connect_mock = MagicMock(return_value=1)
    monkeypatch.setattr(connector, 'connect', connect_mock)
    assert SnowflakeSQLExecutor(sf_config=sf_config_mock).get_connection() == 1


def test_get_cursor_autocommit(monkeypatch, sf_config_mock):
    """Test get_cursor method in autocommit mode."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        with sf_exec.get_cursor() as cursor:
            pass
    cursor.assert_has_calls([call.close()])


def test_get_cursor_autocommit_disabled(monkeypatch, sf_config_mock):
    """Test get_cursor method in non-autocommit mode."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(
            sf_config=sf_config_mock, autocommit=False) as sf_exec:
        with sf_exec.get_cursor() as cursor:
            pass
    cursor.assert_has_calls([call.execute('BEGIN')], any_order=True)
    connect_mock.assert_has_calls([call().commit()], any_order=True)


def test_execute(monkeypatch, sf_config_mock):
    """Test execute method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.execute('SELECT %(x)s;', {'x': 'y'})
    connect_mock.assert_has_calls([
        call().cursor(),
        call().cursor().execute('SELECT %(x)s;', {'x': 'y'}),
        call().cursor().close()], any_order=True)


def test_executemany(monkeypatch, sf_config_mock):
    """Test executemany method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.executemany('SELECT %(x)s;', [{'x': 'y'}, {'x': 'z'}])
    connect_mock.assert_has_calls([
        call().cursor(),
        call().cursor().executemany('SELECT %(x)s;', [{'x': 'y'}, {'x': 'z'}]),
        call().cursor().close()], any_order=True)


def test_fetchone(monkeypatch, sf_config_mock):
    """Test fetchone method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(
            sf_config=sf_config_mock) as sf_exec:
        sf_exec.fetchone('SELECT %(x)s;', {'x': 'y'})
    connect_mock.assert_has_calls([
        call().cursor(),
        call().cursor().execute('SELECT %(x)s;', {'x': 'y'}),
        call().cursor().fetchone(),
        call().cursor().close()], any_order=True)


def test_fetchall(monkeypatch, sf_config_mock):
    """Test fetchall method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.fetchall('SELECT %(x)s;', {'x': 'y'})
    connect_mock.assert_has_calls([
        call().cursor(),
        call().cursor().execute('SELECT %(x)s;', {'x': 'y'}),
        call().cursor().fetchall(),
        call().cursor().close()], any_order=True)


def test_fetchmany(monkeypatch, sf_config_mock):
    """Test fetchmany method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        gen = sf_exec.fetchmany('SELECT %(x)s;', 5, {'x': 'y'})
    next(gen)
    connect_mock.assert_has_calls([
        call().cursor(),
        call().cursor().execute('SELECT %(x)s;', {'x': 'y'}),
        call().cursor().fetchmany(5)], any_order=True)


def test_table_exists(monkeypatch, sf_config_mock):
    """Test table_exists method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.table_exists('some_table')
    connect_mock.assert_has_calls([
        call().cursor().execute(
            'SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES '
            'WHERE TABLE_NAME = %(table)s AND TABLE_SCHEMA = %(schema)s '
            'AND TABLE_CATALOG = %(db)s;',
            {'db': 'TEST_DB', 'schema': 'TEST_SCHEMA', 'table': 'SOME_TABLE'}),
        call().cursor().fetchone()])


def test_drop_table(monkeypatch, sf_config_mock):
    """Test drop_table method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.drop_table('some_table')
    connect_mock.assert_has_calls([
        call().cursor().execute(
            'DROP TABLE IF EXISTS test_db.test_schema.some_table;')],
        any_order=True)


def test_truncate_table(monkeypatch, sf_config_mock):
    """Test truncate_table method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.truncate_table('some_table')
    connect_mock.assert_has_calls([
        call().cursor().execute(
            'TRUNCATE test_db.test_schema.some_table;')], any_order=True)


def test_get_column_names(monkeypatch, sf_config_mock):
    """Test get_column_names method."""
    connect_mock = MagicMock()
    connect_mock.return_value.cursor.return_value.fetchall.return_value = (
        ('c1', 'desc'), ('c2', 'desc'))
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        column_names = sf_exec.get_column_names('some_table')
    connect_mock.assert_has_calls([
        call().cursor().execute(
            'DESC TABLE test_db.test_schema.some_table;', None)],
        any_order=True)
    assert column_names == ['c1', 'c2']


def test_swap_tables(monkeypatch, sf_config_mock):
    """Test swap_tables method."""
    connect_mock = MagicMock()
    monkeypatch.setattr(connector, 'connect', connect_mock)
    with SnowflakeSQLExecutor(sf_config=sf_config_mock) as sf_exec:
        sf_exec.swap_tables('table1', 'table2')
    connect_mock.assert_has_calls([
        call().cursor().execute(
            'ALTER TABLE test_db.test_schema.table1 SWAP WITH '
            'test_db.test_schema.table2;')], any_order=True)