"""Utilities for working with DB in tests.""" from contextlib import contextmanager from functools import wraps import sys from unittest import mock from tests.unit.factories.alw_user_vend_contact import ALWUserVendContactFactory # noqa from tests.unit.factories.subaccount import SubaccountFactory from tests.unit.factories.vendor import VendorFactory from users import config from users.connectors import mysql from users.connectors.mysql import db_session def seed_models(models): """Save the given model(s) to the DB. Args: models (list): list of model instances to save. """ if not hasattr(models, '__iter__'): models = [models] with mysql.db_session() as session: _exit_if_not_test_environment(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_vw_active_vendor_contract(): """Create the `vw_active_vendor_contract` view.""" with db_session() as session: session.execute('DROP VIEW IF EXISTS vw_active_vendor_contract') session.execute( 'CREATE VIEW vw_active_vendor_contract as \ SELECT v.vendor_id, vc.id as vendor_contract_id \ FROM vendor v \ INNER JOIN vendor_contract vc ON v.vendor_id = vc.id' ) 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.db_session() as session: _exit_if_not_test_environment(session) mysql.BaseModel.metadata.create_all(mysql.db_engine) create_vw_active_vendor_contract() try: function_return = function(*args, **kwargs) finally: mysql.BaseModel.metadata.drop_all(mysql.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.TEST_ENVIRONMENT: sys.exit('Environment must be set to {}.'.format(config.TEST_ENVIRONMENT)) 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, 'db_session', fake_session_manager) return mock_session def seed_vendor_vc_subaccount(data): """Seed vendor, vend_contact and subaccount tables.""" result_data = [] all_vendors = {} for user in data: if user.get('vendor_id') not in all_vendors: vendor = VendorFactory.build( vendor_id=user.get('vendor_id'), company=user.get('company'), support_contact_email=user.get('support_contact_email', ''), ) seed_models(vendor) all_vendors[user.get('vendor_id')] = vendor else: vendor = all_vendors[user.get('vendor_id')] vend_contact = ALWUserVendContactFactory.build( user_id=user.get('id'), vendor_id=user.get('vendor_id'), auth0_user_id=user.get('auth0_id'), active=user.get('active', 'Y'), subaccount_id=user.get('subaccount_id'), ) seed_models(vend_contact) subaccount_data = {} if user.get('subaccount_id'): subaccount = SubaccountFactory.build(subaccount_id=user.get('subaccount_id')) seed_models(subaccount) subaccount_data = { 'subaccount_id': subaccount.subaccount_id, 'subaccount_name': subaccount.subaccount_name, } result_data.append( { 'vendor_name': vendor.name, 'vendor_id': vendor.vendor_id, 'subaccount_name': subaccount_data.get('subaccount_name'), 'subaccount_id': subaccount_data.get('subaccount_id'), 'vc_id': vend_contact.user_id, 'auth0_user_id': vend_contact.auth0_user_id, 'auth0_primary': vend_contact.auth0_primary, 'label_identifier': vendor.label_identifier, 'company': vendor.company, } ) return result_data