"""Utility functions to support database interactions in tests."""

import sys
from contextlib import contextmanager
from functools import wraps
from unittest import mock

import common_config as config

import mysql


def test_schema(function):
    """Test schema.

    Decorator that creates the test DB schema before a function call and
    tears the schema down after the function call has finished.
    This just creates the schema and does not seed data. Indvidual test cases
    can use factories to seed data as needed.

    Args:
        Function (func): function to be called after creating the test schema.
    Returns:
        Function: The decorated function.
    """
    @wraps(function)
    def call_function_within_db_context(*args, **kwargs):
        with mysql.ar_db_session() as session:
            _exit_if_not_test_environment(session)
        mysql.BaseModel.metadata.create_all(mysql.ar_db_engine)
        try:
            function_return = function(*args, **kwargs)
        finally:
            mysql.BaseModel.metadata.drop_all(mysql.ar_db_engine)

        return function_return
    return call_function_within_db_context


def _exit_if_not_test_environment(session):
    """For safety, only run tests in test environment pointed to sqlite.

    Exit immediately if not in test environment or not pointed to sqlite.
    """
    if config.ENVIRONMENT != config.ENVIRONMENT_TEST:
        sys.exit('Environment must be set to {}.'.format(
            config.ENVIRONMENT_TEST))
    if 'sqlite' not in session.bind.url.drivername:
        sys.exit('Tests must point to sqlite database.')


def mock_db_session(mocker):
    """Create a mock database session.

    Also mock the db_session context manager to use the mock session.
    """
    mock_session = mock.Mock(query=mock.Mock())

    @contextmanager
    def fake_session_manager():
        yield mock_session

    mocker.patch.object(mysql, 'ar_db_session', fake_session_manager)

    return mock_session