"""Database test utilities."""

from functools import wraps
import sys

from collaborator import config
from collaborator.connectors import snowflake


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.TEST_ENVIRONMENT:
        sys.exit(f"Environment must be set to {config.TEST_ENVIRONMENT}.")

    if "sqlite" not in session.bind.url.drivername:
        sys.exit("Tests must point to sqlite database.")


@snowflake.db_session
def _create_tables(session):
    """Create the split tables."""
    _exit_if_not_test_environment(session)
    snowflake.BaseModel.metadata.create_all(snowflake._db_engine)


@snowflake.db_session
def _drop_tables(session):
    """Drop all tables."""
    _exit_if_not_test_environment(session)
    snowflake.BaseModel.metadata.drop_all(snowflake._db_engine)


def test_schema_default_seed(function):
    """Create and tear down the test DB schema around a function call.

    Args:
        Function (func): the function to be called after creating the test
        schema.

    Returns:
        Function: The decorated function.
    """

    @wraps(function)
    def wrapper(*args, **kwargs):
        _create_tables()

        try:
            function_return = function(*args, **kwargs)
        finally:
            _drop_tables()

        return function_return

    return wrapper