"""Utility functions.""" from datetime import datetime from datetime import timedelta from functools import wraps from salessheets.connectors import mysql def _subdate(date, days): """Subtract days from given date. This function is used as replacement for MySQL function SUBDATE. Since SQLite does not have such function, but it's used in application. To make tests running on SQLite in-memory database, this function is registered as user-defined. Args: date (str): Date as formatted string. This is how date is stored in SQLite by sqlalchemy. days (int): Number of days to subtract from date. Returns: str: Result date as formatted string. """ fmt = '%Y-%m-%d %H:%M:%S' date = datetime.strptime(date, fmt) new_date = date - timedelta(days=days) return new_date.strftime(fmt) def init_database(): """Create all tables in the configured database. Expects that no table exist. Useful only for tests setup functions. """ mysql.BaseModel.metadata.create_all(mysql._salessheets_history_engine) if mysql._salessheets_history_engine.name == ( 'sqlite'): # pragma: no cover with mysql._salessheets_history_engine.begin() as sa_conn: sa_conn.connection.connection.create_function( 'SUBDATE', 2, _subdate) def db_setup_function(): """Setup to perform before running any test function in this module. Create all tables in the test (SQLite) database. Import as 'setup_function' or use directly in test modules. """ init_database() def db_teardown_function(): """Tear down to perform after running any test function in this module. Drop test (SQLite) database. Import as 'setup_function' or use directly in test modules. """ mysql.BaseModel.metadata.drop_all(mysql._salessheets_history_engine) def seed_ss_models(models): """Save the given model(s) to the DB. Args: models (list): list of model instances to save. """ models = list(models) with mysql.salessheets_history_session_scope() as session: for model in models: session.add(model) session.flush() # detach the objects from this session so tests can interrogate them for model in models: session.expunge(model) def create_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.salessheets_history_session_scope(): mysql.BaseModel.metadata.create_all( mysql._salessheets_history_engine) try: function_return = function(*args, **kwargs) finally: mysql.BaseModel.metadata.drop_all( mysql._salessheets_history_engine) return function_return return call_function_within_db_context