"""Tests for analytics_aggregation.util.db module."""

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

from pymysql.connections import Connection
from pymysql.cursors import Cursor
import pytest
from pytest import raises

import analytics_aggregation.util.db as db


@pytest.fixture()
def redshift_config():
    """Fixture with dummy redshift configuration."""
    return {
        'driver': db.Driver.REDSHIFT, 'host': 'host', 'port': 1234,
        'user': 'user', 'password': 'password', 'db': 'db'}


@pytest.fixture()
def mysql_config():
    """Fixture with dummy mysql configuration."""
    return {
        'driver': db.Driver.MYSQL, 'host': 'host', 'user': 'user',
        'password': 'password', 'db': 'db'}


def mock_db_context(db_context):
    """Mock db_context."""
    cursor = Mock(spec=Cursor)
    connection = Mock(spec=Connection)

    context = Mock()
    context.__enter__ = Mock(return_value=(cursor, connection))
    context.__exit__ = Mock()

    db_context.return_value = context
    db_context.mocked_cursor = cursor
    db_context.mocked_connection = connection
    return db_context


@patch('analytics_aggregation.util.db.psycopg2')
def test_connect_redshift(psycopg2, redshift_config):
    """Test connect function."""
    connect = psycopg2.connect = MagicMock()
    db.connect(redshift_config)
    connect.assert_called_once_with(
        host=redshift_config['host'],
        port=redshift_config['port'],
        user=redshift_config['user'],
        password=redshift_config['password'],
        database=redshift_config['db'])


@patch('analytics_aggregation.util.db.pymysql')
def test_connect_mysql(psycopg2, mysql_config):
    """Test connect function."""
    connect = psycopg2.connect = MagicMock()
    db.connect(mysql_config)
    connect.assert_called_once_with(
        host=mysql_config['host'],
        user=mysql_config['user'],
        password=mysql_config['password'],
        db=mysql_config['db'])


def test_connect_not_implemented():
    """Test connect function for unsupported DB driver."""
    db_conf = {
        'driver': 'unknown'}
    with raises(db.UnsupportedDriverException):
        db.connect(db_conf)


def test_query_fill():
    """Test fill method from the Query class."""
    query_string = 'SELECT * FROM {table_name};'
    table_name = 'some_table'
    expected_query_string = 'SELECT * FROM some_table;'

    query = db.Query(query_string)
    assert query.sql == query_string

    query.fill(table_name=table_name)
    assert query.sql == expected_query_string


@patch('analytics_aggregation.util.db.connect')
@patch('analytics_aggregation.util.db.db_context')
def test_query_execute(db_context, connect):
    """Test execute method from the Query class."""
    db_context = mock_db_context(db_context)
    query_string = 'SQL'
    query = db.Query(query_string)

    result = query.execute({})
    db_context.mocked_cursor.execute.assert_called_once_with(query_string, {})

    assert result is None


@patch('analytics_aggregation.util.db.connect')
@patch('analytics_aggregation.util.db.db_context')
def test_query_fetchall(db_context, connect):
    """Test fetchall method from the Query class."""
    db_context = mock_db_context(db_context)
    query_result = [(1, 2, 3), (3, 4, 5)]
    db_context.mocked_cursor.fetchall.return_value = query_result

    query_string = 'SQL'
    query = db.Query(query_string)

    result = query.fetchall({})
    assert result == query_result
    db_context.mocked_cursor.execute.assert_called_once_with(query_string, {})
    db_context.mocked_cursor.fetchall.assert_called_once_with()


@patch('analytics_aggregation.util.db.connect')
@patch('analytics_aggregation.util.db.db_context')
def test_query_fetchone(db_context, connect):
    """Test fetchone method from the Query class."""
    db_context = mock_db_context(db_context)
    query_result = [(1, 2, 3)]
    db_context.mocked_cursor.fetchone.return_value = query_result

    query_string = 'SQL'
    query = db.Query(query_string)

    result = query.fetchone({})
    assert result == query_result
    db_context.mocked_cursor.execute.assert_called_once_with(query_string, {})
    db_context.mocked_cursor.fetchone.assert_called_once_with()


@patch('analytics_aggregation.util.db.connect')
@patch('analytics_aggregation.util.db.db_context')
def test_query_fetchmany(db_context, connect):
    """Test fetchmany generator from the Query class."""
    db_context = mock_db_context(db_context)
    db_context.mocked_cursor.fetchmany.return_value = [(1, 2), (2, 3)]
    query_string = 'SQL'
    query = db.Query(query_string)

    gen = query.fetchmany({}, 2)
    next(gen)
    db_context.mocked_cursor.execute.assert_called_once_with(query_string, {})
    db_context.mocked_cursor.fetchmany.assert_called_once_with(2)


@patch('analytics_aggregation.util.db.connect')
@patch('analytics_aggregation.util.db.db_context')
def test_query_fetchmany_empty_result(db_context, connect):
    """Test fetchmany generator from the Query class."""
    db_context = mock_db_context(db_context)
    db_context.mocked_cursor.fetchmany.return_value = None
    query_string = 'SQL'
    query = db.Query(query_string)
    gen = query.fetchmany({}, 2)
    with raises(StopIteration):
        next(gen)


def test_sqlloader_constructor():
    """Test for SQLLoader constructor."""
    query_root = 'path'
    sql = db.SQLLoader(query_root)
    assert sql.query_cash == {}
    assert sql.sql_files_root == query_root


@patch('builtins.open')
def test_sqlloader_get_item(open_mock):
    """Test getting SQL queries from disk."""
    q_string = 'SELECT * FROM {somewhere};'
    open_mock.return_value.__enter__.return_value.read.return_value = q_string

    sql = db.SQLLoader('path')
    query = sql['query']
    assert query.sql == q_string
    open_mock.assert_called_once_with('path/query.sql', 'r')


@patch('builtins.open')
def test_sqlloader_get_item_from_cache(open_mock):
    """Test getting SQL queries from cache."""
    q_string = 'SELECT * FROM {somewhere};'
    sql = db.SQLLoader('path')
    sql.query_cash = {'query': q_string}

    query = sql['query']
    assert query.sql == q_string
    open_mock.assert_not_called()


def test_db_context():
    """Test db_context function."""
    cursor = Mock()
    connection = Mock()
    connection.cursor.return_value = cursor

    with db.db_context(connection) as (tcursor, tconnection):
        # these asserts also guarantee rollbacks and explicit commits are
        # called
        assert tcursor == cursor
        assert tconnection == connection

        tcursor.execute('foo')
        assert not tconnection.commit.called

    cursor.execute.assert_called_once_with('foo')
    assert cursor.close.called
    assert connection.commit.called
    assert connection.close.called

    # rollback on exception
    cursor.reset_mock()
    connection.reset_mock()

    with raises(BaseException):
        with db.db_context(connection) as (tcursor, tconnection):
            tcursor.execute('foo bar baz')
            raise BaseException()

    assert not connection.commit.called
    cursor.execute.assert_called_once_with('foo bar baz')
    assert connection.rollback.called


def test_query_fill_escaped_with_mixed_parameters():
    """Test fill_escaped fails when passing mixed parameters."""
    query_string = 'SQL %s %(keyword)s'
    query = db.Query(query_string)
    with raises(ValueError):
        query.fill_escaped('1', keyword='2')


def test_query_fill_escaped_with_positional_parameters():
    """Test fill_escaped work with positional arguments."""
    query_string = 'SQL %s'
    query = db.Query(query_string)
    query.fill_escaped('1')
    assert type(query.escape_parameters) == tuple
    assert query.escape_parameters == ('1', )


def test_query_fill_escaped_with_keyword_parameters():
    """Test fill_escaped work with keyword arguments."""
    query_string = 'SQL %(keyword)s'
    query = db.Query(query_string)
    query.fill_escaped(keyword='1')
    assert type(query.escape_parameters) == dict
    assert query.escape_parameters == {'keyword': '1'}


def test_escape_quotes():
    """Test escape_quotes function."""
    expected = r"\'hello\'"
    initial = "'hello'"
    assert db.escape_quotes(initial) == expected