"""Utility functions to support database interactions in tests.""" from contextlib import contextmanager from functools import wraps import sys from unittest import mock from product_workflow import config from product_workflow.connectors import mysql from product_workflow.models.meta_update_queue import MetaUpdateDmsMaster from product_workflow.models.meta_update_queue import MetaUpdateQueue CREATE_RELEASE_APPROVAL_QUEUE_TABLE = """ CREATE TABLE IF NOT EXISTS `release_approval_queue` ( `release_approval_id` INTEGER NOT NULL PRIMARY KEY, `release_id` INTEGER NOT NULL, `release_correction_id` INTEGER, `status` TEXT NOT NULL, `checked_out_by` INTEGER, `admin_approval` VARCHAR(1), `last_updated` DATETIME NOT NULL, `date_submitted` DATETIME NOT NULL ) """ DROP_RELEASE_APPROVAL_QUEUE_TABLE = """ DROP TABLE IF EXISTS `release_approval_queue` """ CREATE_REJECTION_NOTES_TABLE = """ CREATE TABLE IF NOT EXISTS `rejection_notes` ( `rejection_id` INTEGER NOT NULL PRIMARY KEY, `release_approval_id` INTEGER NOT NULL, `table_name` VARCHAR(45), `field_name` VARCHAR(45), `comments` TEXT, `corrected` VARCHAR(1), `date_added` DATETIME NOT NULL, `key_id` INTEGER ) """ DROP_REJECTION_NOTES_TABLE = """ DROP TABLE IF EXISTS `rejection_notes` """ CREATE_RELEASE_CORRECTION_TABLE = """ CREATE TABLE IF NOT EXISTS `release_correction` ( `release_correction_id` INTEGER NOT NULL PRIMARY KEY, `release_id` int(10) NOT NULL, `status` VARCHAR(10) NOT NULL, `last_updated` datetime NOT NULL, `last_updated_by` int(10) NOT NULL, `last_updated_type` VARCHAR(10) DEFAULT 'vendor' ) """ DROP_RELEASE_CORRECTION_TABLE = """ DROP TABLE IF EXISTS `release_correction` """ CREATE_RELEASE_CORRECTION_DETAIL_TABLE = """ CREATE TABLE IF NOT EXISTS `release_correction_detail` ( `release_correction_detail_id` INTEGER NOT NULL PRIMARY KEY, `release_correction_id` int(10) DEFAULT NULL, `table_name` varchar(45) NOT NULL, `field_name` varchar(45) DEFAULT NULL, `key_id` int(11) DEFAULT NULL, `key_value` text NOT NULL, `last_updated` datetime NOT NULL, `last_updated_by` int(10) NOT NULL, `last_updated_type` VARCHAR(10) DEFAULT NULL ) """ CREATE_META_UPDATE_QUEUE_TABLE = """ CREATE TABLE IF NOT EXISTS `meta_update_queue` ( `meta_update_queue_id` int(11), `upc` bigint(20) DEFAULT NULL, `orchadmin_user_id` int(11) DEFAULT NULL, `date_added` datetime DEFAULT NULL, `update_type` DEFAULT NULL, `description` text, `status` DEFAULT 'new' ) """ CREATE_META_UPDATE_DMS_MASTER_TABLE = """ CREATE TABLE IF NOT EXISTS `meta_update_dms_master` ( `meta_update_dms_master_id` int(10), `meta_update_queue_id` int(10), `customer_master_master_id` smallint(5) NOT NULL, `date_processed` datetime DEFAULT NULL ) """ DROP_RELEASE_CORRECTION_DETAIL_TABLE = """ DROP TABLE IF EXISTS `release_correction_detail` """ DROP_META_UPDATE_QUEUE_TABLE = """ DROP TABLE IF EXISTS `meta_update_queue` """ DROP_META_UPDATE_DMS_MASTER_TABLE = """ DROP TABLE IF EXISTS `meta_update_dms_master` """ def execute_sql(*args): """Run a series of queries safely within the test environment.""" with mysql.db_session() as session: _exit_if_not_test_environment(session) for query in args: session.execute(query) 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 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): execute_sql( CREATE_RELEASE_APPROVAL_QUEUE_TABLE, CREATE_REJECTION_NOTES_TABLE, CREATE_RELEASE_CORRECTION_TABLE, CREATE_RELEASE_CORRECTION_DETAIL_TABLE, CREATE_META_UPDATE_QUEUE_TABLE, CREATE_META_UPDATE_DMS_MASTER_TABLE) try: function_return = function(*args, **kwargs) finally: execute_sql( DROP_RELEASE_APPROVAL_QUEUE_TABLE, DROP_REJECTION_NOTES_TABLE, DROP_RELEASE_CORRECTION_TABLE, DROP_RELEASE_CORRECTION_DETAIL_TABLE, DROP_META_UPDATE_QUEUE_TABLE, DROP_META_UPDATE_DMS_MASTER_TABLE) 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 sesssion. 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 insert_into_meta_update_queue(): """Insert data into scheduled_product_update and associated tables.""" upc = 123 with mysql.db_session() as session: new_record = MetaUpdateQueue( meta_update_queue_id=123, upc=upc, update=list(map( lambda x: MetaUpdateDmsMaster( customer_master_master_id=x ), [1, 2, 3] )) ) session.add(new_record) session.commit() return new_record.to_dict()