"""db.py. Database level utility class for testing against timed_release model. """ import datetime import sys from functools import wraps from timed_release import config from timed_release.connectors.sql import ( ar_db_engine, ar_db_session_maker, base_model, db_engine, db_session, db_session_maker, ) from timed_release.models import ( product_live_time, product_store, timed_release, ) from timed_release.models.scheduled_update import ( ScheduledProductUpdate, ScheduledProductUpdateDms, ScheduledProductUpdateValues, ) def create_all_tables(): """Create database table for all model classes.""" _exit_if_not_test_environment(db_session_maker()) _exit_if_not_test_environment(ar_db_session_maker()) # Creates table for all model where base_model is inherited. base_model.metadata.create_all(db_engine) base_model.metadata.create_all(ar_db_engine) def drop_all_tables(): """Drop a database table for model classes.""" _exit_if_not_test_environment(db_session_maker()) # Drops table for all model where base_model is inherited. base_model.metadata.drop_all(db_engine) 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 test_schema(function): """Create and tear down the test DB schema around a function call. This just creates the schema and does not seed data. Individual test cases can use factories to seed data as needed. Args: function (func): the function to be called after creating the test schema. Returns: Function: The decorated function. """ @wraps(function) def call_function_within_db_context(*args, **kwargs): create_all_tables() try: function_return = function(*args, **kwargs) finally: drop_all_tables() return function_return return call_function_within_db_context INSERT_DATA = [ { 'product_id': '12', 'time_of_day_product': datetime.time(20, 30, 00), 'time_zone': 'GMT' }] def insert_product_live_time_data(): """Insert data to product_live_time table.""" with db_session() as session: session.bulk_insert_mappings( product_live_time.ProductLiveTime, INSERT_DATA) INSERT_DATA_PRODUCT_STORE = [ { 'product_id': '12', 'store_id': 1 }, { 'product_id': '12', 'store_id': 2 }] def insert_product_store_data(): """Insert data to product_store table.""" with db_session() as session: session.bulk_insert_mappings( product_store.ProductStore, INSERT_DATA_PRODUCT_STORE) INSERT_DATA_TIMED_RELEASE = [ { 'product_id': 12, 'timing': 'staggered', 'store_id': 1, 'sales_date_time': datetime.datetime(2023, 3, 3, 0, 0, 0), 'created_on': datetime.datetime(2023, 3, 3, 10, 10, 10), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2023, 3, 3, 10, 10, 10), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'source': 'oa', 'user_timezone': 'america/new_york' }, { 'product_id': 12, 'timing': 'timed', 'store_id': 286, 'sales_date_time': datetime.datetime(2023, 3, 3, 10, 10, 10), 'created_on': datetime.datetime(2023, 3, 3, 10, 10, 10), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2023, 3, 3, 10, 10, 10), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'source': 'oa', 'user_timezone': 'america/new_york' }, { 'product_id': 13, 'timing': 'timed', 'store_id': 187, 'sales_date_time': datetime.datetime(2026, 5, 20, 11, 11, 11), 'created_on': datetime.datetime(2025, 8, 14, 9, 20, 15), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2025, 8, 14, 9, 20, 15), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'source': 'workstation', 'user_timezone': 'america/new_york' }, { 'product_id': 13, 'timing': 'timed', 'store_id': 286, 'sales_date_time': datetime.datetime(2026, 5, 20, 11, 11, 11), 'created_on': datetime.datetime(2025, 8, 14, 9, 20, 15), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2025, 8, 14, 9, 20, 15), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'source': 'workstation', 'user_timezone': 'america/new_york' }] def insert_timed_release_data(): """Insert data to timed_release table.""" with db_session() as session: session.bulk_insert_mappings( timed_release.TimedRelease, INSERT_DATA_TIMED_RELEASE) def insert_scheduled_update_data(): """Insert data into scheduled_product_update and associated tables.""" orchard_identity_id = 'b211c821-8b5f-40bd-91a6-d139fb9a95f6' with db_session() as session: schedules = [ ScheduledProductUpdate( product_id=123, sales_date_time=datetime.datetime(2023, 3, 2, 0, 0), process_at=datetime.datetime(2023, 3, 2, 0, 0), process_at_offset='00:00', scheduled_by=orchard_identity_id, created_on=datetime.datetime(2023, 3, 3), created_by=orchard_identity_id, updated_on=datetime.datetime(2023, 3, 3), updated_by=orchard_identity_id, dsps=([ ScheduledProductUpdateDms(customer_master_master_id=x) for x in [1, 2, 3] ]), update=ScheduledProductUpdateValues( sale_start_date=datetime.datetime(2023, 3, 3) ), source='oa', user_timezone='america/new_york' ), ScheduledProductUpdate( product_id=124, sales_date_time=datetime.datetime(2025, 9, 24, 0, 0), process_at=datetime.datetime(2025, 9, 24, 0, 0), process_at_offset='00:00', scheduled_by=orchard_identity_id, created_on=datetime.datetime(2023, 8, 14), created_by=orchard_identity_id, updated_on=datetime.datetime(2023, 8, 14), updated_by=orchard_identity_id, user_timezone='UTC', dsps=([ ScheduledProductUpdateDms(customer_master_master_id=x) for x in [1, 2, 3] ]), update=ScheduledProductUpdateValues( sale_start_date=datetime.datetime(2025, 9, 24) ), source='workstation' ), ScheduledProductUpdate( product_id=125, sales_date_time=datetime.datetime(2035, 3, 2, 0, 0), process_at=datetime.datetime(2035, 3, 2, 0, 0), process_at_offset='00:00', scheduled_by=orchard_identity_id, created_on=datetime.datetime(2035, 3, 3), created_by=orchard_identity_id, updated_on=datetime.datetime(2035, 3, 3), updated_by=orchard_identity_id, dsps=([ ScheduledProductUpdateDms(customer_master_master_id=x) for x in [4, 5, 6] ]), update=ScheduledProductUpdateValues( sale_start_date=datetime.datetime(2035, 3, 3) ), source='oa' ), ScheduledProductUpdate( product_id=126, sales_date_time=datetime.datetime(2035, 9, 24, 18, 0), process_at=datetime.datetime(2035, 9, 24, 17, 30), process_at_offset='00:00', scheduled_by=orchard_identity_id, created_on=datetime.datetime(2035, 8, 14), created_by=orchard_identity_id, updated_on=datetime.datetime(2035, 8, 14), updated_by=orchard_identity_id, dsps=([ ScheduledProductUpdateDms(customer_master_master_id=x) for x in [7, 8, 9] ]), update=ScheduledProductUpdateValues( sale_start_date=datetime.datetime(2035, 9, 24) ), source='workstation', user_timezone='america/new_york' ) ] session.add_all(schedules) session.commit() return [schedule.to_dict() for schedule in schedules] INSERT_DATA_WORKSTATION_TIMED_RELEASE = [ { 'product_id': 1, 'sales_start_datetime': datetime.datetime(2025, 4, 3, 10, 00, 00), 'event_process_on': datetime.datetime(2025, 4, 3, 9, 30, 00), 'created_on': datetime.datetime(2025, 4, 1, 10, 10, 10), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2025, 4, 1, 10, 10, 10), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'is_deleted': False }, { 'product_id': 2, 'sales_start_datetime': datetime.datetime(2025, 6, 12, 9, 0, 0), 'event_process_on': datetime.datetime(2025, 6, 12, 8, 30, 0), 'created_on': datetime.datetime(2025, 4, 22, 10, 10, 10), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2025, 4, 22, 10, 10, 10), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'is_deleted': False }, { 'product_id': 3, 'sales_start_datetime': datetime.datetime(2030, 8, 12, 5, 0, 0), 'event_process_on': datetime.datetime(2030, 8, 12, 4, 30, 0), 'created_on': datetime.datetime(2025, 6, 29, 9, 12, 0), 'created_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'updated_on': datetime.datetime(2025, 6, 29, 9, 12, 0), 'updated_by': '03e83149-3c7f-4587-940b-ce83dec34bc8', 'is_deleted': True } ]