import os import unittest from datetime import datetime, timedelta, timezone from threading import Thread from time import sleep from unittest.mock import Mock import pytest from parameterized import parameterized from db_schema.postgres import connection from db_schema.schemas.exploration import DisassembleContentStatus, DisassembleStatusEnum from db_schema.schemas.slz import ( ActivityStatusEnum, CompletenessStatusEnum, ContentFailureLog, ContentMetadataStatusEnum, ContentStatus, ContentStatusEnum, ContentStatusMigrationLog, DataSource, Licensor, Report, UnitOfWork, UnitOfWorkMigrationLog, UnitOfWorkPriorityEnum ) from psycopg2.extras import DateTimeRange from sqlalchemy import text from sqlalchemy.orm import Session from exp_storage.conv import UoWPGToReadable from exp_storage.repository import Repository from tests.utils.stub import content_status_stub, disassemble_cs_stub, uow_stub @pytest.mark.integration class DisassembleContentStatusTestCase(unittest.TestCase): def setUp(self): logger = Mock() self._pg = connection.get_session( host=os.environ.get('PG_HOST', '0.0.0.0'), port=os.environ.get('PG_PORT', 5432), db=os.environ.get('PG_DB', 'slz'), user=os.environ.get('PG_USER', 'admin'), password=os.environ.get('PG_PASSWORD', 'admin'), engine_params={'echo': True} ) self._repository = Repository(logger, self._pg) self._pg = self._repository._session self._licensors = {item.licensor_name: item for item in self._pg.query(Licensor).all()} self._reports = {item.report_name: item for item in self._pg.query(Report).all()} self._data_source = { item.data_source_name: item for item in self._pg.query(DataSource).all() } now = datetime.now(tz=timezone.utc) self._now = now self._uows = { 1: dict( unit_of_work_code='apple-20191124-theorchard-amEvent-v1_2', reprocess_id='1', licensor=self._licensors['theorchard'], report=self._reports['amEvent'], report_date='2019-11-24', version='v1_2', activity_status=ActivityStatusEnum.NOT_IN_PROGRESS.value, completeness_status=CompletenessStatusEnum.COMPLETE.value, is_force_complete=False, priority=UnitOfWorkPriorityEnum.DEFAULT.value, next_run_at=now + timedelta(minutes=15), created_at=now, last_updated_at=now, ), 2: dict( unit_of_work_code='apple-20191125-theorchard-amEvent-v1_2', reprocess_id='1', licensor=self._licensors['theorchard'], report=self._reports['amEvent'], report_date='2019-11-25', version='v1_2', activity_status=ActivityStatusEnum.NOT_IN_PROGRESS.value, completeness_status=CompletenessStatusEnum.COMPLETE.value, is_force_complete=False, priority=UnitOfWorkPriorityEnum.DEFAULT.value, next_run_at=now + timedelta(minutes=15), created_at=now, last_updated_at=now, ), } self._content_statuses = { 1: dict( context='US', content_name='us.txt', content_status=ContentStatusEnum.COMPLETE, created_at=now, failure_count=0, sub_content='{}', metadata_process_status=ContentMetadataStatusEnum.NOT_QUEUED.value, ), } self._pg.rollback() def tearDown(self): # clean up tables self._pg.rollback() self._pg.query(UnitOfWorkMigrationLog).delete() self._pg.query(ContentStatusMigrationLog).delete() self._pg.query(ContentFailureLog).delete() self._pg.query(DisassembleContentStatus).delete() self._pg.query(ContentStatus).delete() self._pg.query(UnitOfWork).delete() self._pg.commit() self._pg.close() def test_get_disassemble_cs(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) status = DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() result = self._repository.get_disassemble_cs( uow_id=UoWPGToReadable.get_uow_id(uow_1), context=cs_1.context, content_name="test_content_name" ) self.assertTrue(result, status) def test_get_disassemble_cs_by_id(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) status = DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() actual = self._pg.query(DisassembleContentStatus).first() result = self._repository.get_disassemble_cs_by_id( disassemble_cs_id=actual.disassemble_content_status_id ) self.assertTrue(result, actual) def test_create_disassemble_content_status(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) self._pg.commit() self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 0) changeset = { 'unit_of_work_id': uow_1.unit_of_work_id, 'content_name': "test_content_name", 'context': 'US', 'payload': { 'test': 'test' }, 'dt_now': self._now, } result = self._repository.create_disassemble_content_status(**changeset) self.assertTrue(result) created: DisassembleContentStatus = self._pg.query(DisassembleContentStatus).all()[0] self.assertEqual(created.unit_of_work, uow_1) self.assertEqual(created.content_status, cs_1) self.assertEqual(created.content_name, changeset['content_name']) self.assertEqual(created.created_at, changeset['dt_now']) self.assertEqual(created.last_updated_at, changeset['dt_now']) self.assertEqual(created.disassemble_status, DisassembleStatusEnum.QUEUED) # create one more self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 1) now = datetime.now(tz=timezone.utc) changeset['dt_now'] = now result = self._repository.create_disassemble_content_status(**changeset) self.assertTrue(result) self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 2) updated = self._pg.query(DisassembleContentStatus).all()[1] self.assertEqual(updated.created_at, changeset['dt_now']) self.assertEqual(updated.last_updated_at, changeset['dt_now']) def test_create_disassemble_content_status_by_uow_id(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) self._pg.commit() self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 0) changeset = { 'uow_id': UoWPGToReadable.get_uow_id(uow_1), 'content_name': "test_content_name", 'context': 'US', 'payload': { 'test': 'test' }, 'dt_now': self._now, } result = self._repository.create_disassemble_content_status_by_uow_id(**changeset) self.assertTrue(result) created: DisassembleContentStatus = self._pg.query(DisassembleContentStatus).all()[0] self.assertEqual(created.unit_of_work, uow_1) self.assertEqual(created.content_status, cs_1) self.assertEqual(created.content_name, changeset['content_name']) self.assertEqual(created.created_at, changeset['dt_now']) self.assertEqual(created.last_updated_at, changeset['dt_now']) self.assertEqual(created.disassemble_status, DisassembleStatusEnum.QUEUED) # create one more self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 1) now = datetime.now(tz=timezone.utc) changeset['dt_now'] = now result = self._repository.create_disassemble_content_status_by_uow_id(**changeset) self.assertTrue(result) self.assertEqual(self._pg.query(DisassembleContentStatus).count(), 2) updated = self._pg.query(DisassembleContentStatus).all()[1] self.assertEqual(updated.created_at, changeset['dt_now']) self.assertEqual(updated.last_updated_at, changeset['dt_now']) def test_update_disassemble_status_to_failed(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() rec = self._pg.query(DisassembleContentStatus).one() self.assertFalse(rec.is_failure_transient) changeset = { 'disassemble_cs_id': rec.disassemble_content_status_id, 'process_status': DisassembleStatusEnum.FAILED, 'is_failure_transient': True, } result = self._repository.update_disassemble_status(**changeset) self.assertTrue(result) updated = self._pg.query(DisassembleContentStatus).filter( DisassembleContentStatus.content_status_id == cs_1.content_status_id, DisassembleContentStatus.unit_of_work_id == uow_1.unit_of_work_id, ).one() self.assertEqual(updated.disassemble_status, changeset['process_status']) self.assertEqual(updated.created_at, self._now) self.assertNotEqual(updated.last_updated_at, self._now) self.assertTrue(rec.is_failure_transient) def test_update_disassemble_status_to_complete(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() rec = self._pg.query(DisassembleContentStatus).one() changeset = { 'disassemble_cs_id': rec.disassemble_content_status_id, 'process_status': DisassembleStatusEnum.COMPLETE, 'params': { 'chunks_count': 12, } } self.assertEqual(rec.completed_at, None) self.assertEqual(rec.chunks_count, None) result = self._repository.update_disassemble_status(**changeset) self.assertTrue(result) updated = self._pg.query(DisassembleContentStatus).one() self.assertEqual(updated.disassemble_status, changeset['process_status']) self.assertEqual(updated.created_at, self._now) self.assertNotEqual(updated.last_updated_at, self._now) self.assertNotEqual(updated.completed_at, None) self.assertEqual(updated.chunks_count, 12) self.assertFalse(rec.is_failure_transient) def test_get_queued(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.IN_PROGRESS ) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.FAILED ) self._pg.commit() result = self._repository.get_queued() self.assertEqual(len(result), 1) def test_get_failed_count(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) dcs = disassemble_cs_stub(unit_of_work=uow_1, content_status=cs_1) self._pg.commit() for i in range(4): disassemble_cs_stub( unit_of_work=uow_1, content_status=cs_1, retried_disassemble_content_status_id=dcs.disassemble_content_status_id ) self._pg.commit() count = self._repository.get_retries_count(dcs.disassemble_content_status_id) self.assertEqual(count, 4) def test_get_active_count(self): result = self._repository.get_active_count() self.assertEqual(result, 0) uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.FAILED ) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.IN_PROGRESS ) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.IN_PROGRESS, dbx_run_id=123, ) self._pg.commit() result = self._repository.get_active_count(dbx_only=False) self.assertEqual(result, 2) result = self._repository.get_active_count(dbx_only=True) self.assertEqual(result, 1) def test_fail_stuck(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) now = self._now - timedelta(minutes=90) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=now, last_updated_at=now, disassemble_status=DisassembleStatusEnum.IN_PROGRESS ) self._pg.commit() # check that unit is in place self.assertEqual(self._repository.get_active_count(dbx_only=False), 1) dcs = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(dcs.last_updated_at, now) self.assertEqual(dcs.completed_at, None) # it's not yet time to fail self._repository.fail_stuck(100) self.assertEqual(self._repository.get_active_count(dbx_only=False), 1) dcs = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(dcs.last_updated_at, now) self.assertEqual(dcs.completed_at, None) # should fail self._repository.fail_stuck(20) self.assertEqual(self._repository.get_active_count(dbx_only=False), 0) dcs = self._pg.query(DisassembleContentStatus).one_or_none() self.assertNotEqual(dcs.last_updated_at, now) self.assertEqual(dcs.completed_at, None) def test_update_dbx_run_id_success(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) status = DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.IN_PROGRESS ) self._pg.commit() record = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(record.dbx_run_id, None) dbx_run_id = 123 result = self._repository.update_dbx_run_id( record.disassemble_content_status_id, dbx_run_id ) updated = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(result, True) self.assertEqual(updated.dbx_run_id, dbx_run_id) def test_update_dbx_run_id_not_updated(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) status = DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED # incorrect status ) self._pg.commit() record = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(record.dbx_run_id, None) dbx_run_id = 123 result = self._repository.update_dbx_run_id( record.disassemble_content_status_id, dbx_run_id ) updated = self._pg.query(DisassembleContentStatus).one_or_none() self.assertEqual(result, False) self.assertEqual(updated.dbx_run_id, None) def test_mark_related_records_as_replaced(self): uows = [ uow_stub( unit_of_work_code='apple-20191124-theorchard-amEvent-v1_2' # this is target record ), uow_stub( unit_of_work_code='apple-20191124-theorchard-amEvent-v1_2', # this is target record reprocess_id='12345', ), uow_stub(unit_of_work_code='apple-20191124-sme-amEvent-v1_2', ), uow_stub( unit_of_work_code='apple-20191124-sme-amEvent-v1_2', reprocess_id='12345', ) ] for uow in uows: self._pg.add(uow) content_statuses0 = [ content_status_stub(context='US'), # this is target record content_status_stub(context='AD'), content_status_stub(context='AE'), ] content_statuses1 = [ content_status_stub(context='US'), # this is target record content_status_stub(context='BK'), ] content_statuses2 = [ content_status_stub(context='US'), content_status_stub(context='JP'), ] for cs in content_statuses0: uows[0].content_statuses.append(cs) for cs in content_statuses1: uows[1].content_statuses.append(cs) for cs in content_statuses2: uows[2].content_statuses.append(cs) disassemble_records = [ disassemble_cs_stub( unit_of_work=uows[0], content_status=content_statuses0[0], # this is target record(should be replaced) content_name="US.txt", ), disassemble_cs_stub( unit_of_work=uows[0], content_status=content_statuses0[1], content_name="AD.txt", ), disassemble_cs_stub( unit_of_work=uows[1], content_status=content_statuses1[0], # this is target record content_name="US.txt", ), disassemble_cs_stub( unit_of_work=uows[1], content_status=content_statuses1[1], content_name="BK.txt", ), disassemble_cs_stub( unit_of_work=uows[2], content_status=content_statuses2[0], content_name="US.txt", ), ] self._pg.commit() self.assertEqual( self._pg.query(DisassembleContentStatus).filter( DisassembleContentStatus.disassemble_status == DisassembleStatusEnum.REPLACED ).count(), 0 ) self.assertIsNone(disassemble_records[2].replaced_disassemble_content_status_id) result = self._repository.mark_related_record_as_replaced( record_to_replace=disassemble_records[0], disassemble_cs=disassemble_records[2] ) self.assertEqual(result, True) records = self._pg.query(DisassembleContentStatus).filter( DisassembleContentStatus.disassemble_status == DisassembleStatusEnum.REPLACED ).all() self.assertEqual(len(records), 1) updated = records[0] self.assertIsNotNone(disassemble_records[2].replaced_disassemble_content_status_id) self.assertIsNotNone(updated.replaced_at) def test_mark_related_records_as_replaced_replace_several(self): uows = [ uow_stub( unit_of_work_code='apple-20191124-theorchard-amEvent-v1_2' # this is target record ), uow_stub( unit_of_work_code='apple-20191124-theorchard-amEvent-v1_2', # this is target record reprocess_id='12345', ), uow_stub(unit_of_work_code='apple-20191124-sme-amEvent-v1_2', ) ] for uow in uows: self._pg.add(uow) content_statuses0 = [ content_status_stub(context='US'), # this is target record content_status_stub(context='AD'), content_status_stub(context='AE'), ] content_statuses1 = [ content_status_stub(context='US'), # this is target record content_status_stub(context='BK'), ] for cs in content_statuses0: uows[0].content_statuses.append(cs) for cs in content_statuses1: uows[1].content_statuses.append(cs) disassemble_records = [ disassemble_cs_stub( unit_of_work=uows[0], content_status=content_statuses0[0], content_name="US.txt", ), disassemble_cs_stub( unit_of_work=uows[0], content_status=content_statuses0[0], # this is target record(should be replaced) content_name="US.txt", ), disassemble_cs_stub( unit_of_work=uows[0], content_status=content_statuses0[1], content_name="AD.txt", ), disassemble_cs_stub( unit_of_work=uows[1], content_status=content_statuses1[0], # this is target record content_name="US.txt", ), disassemble_cs_stub( unit_of_work=uows[1], content_status=content_statuses1[1], content_name="BK.txt", ), ] self._pg.commit() self.assertEqual( self._pg.query(DisassembleContentStatus).filter( DisassembleContentStatus.disassemble_status == DisassembleStatusEnum.REPLACED ).count(), 0 ) self.assertIsNone(disassemble_records[2].replaced_disassemble_content_status_id) result = self._repository.mark_related_record_as_replaced( record_to_replace=disassemble_records[0], disassemble_cs=disassemble_records[3] ) self.assertEqual(result, True) records = self._pg.query(DisassembleContentStatus).filter( DisassembleContentStatus.disassemble_status == DisassembleStatusEnum.REPLACED ).all() self.assertEqual(len(records), 1) for updated in records: self.assertIsNotNone(updated.replaced_at) self.assertIsNotNone(disassemble_records[3].replaced_disassemble_content_status_id) def test_advisory_locked_transaction_isolation(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() critical = False def create_session() -> Session: return connection.get_session( host=os.environ.get('PG_HOST', '0.0.0.0'), port=os.environ.get('PG_PORT', 5432), db=os.environ.get('PG_DB', 'slz'), user=os.environ.get('PG_USER', 'admin'), password=os.environ.get('PG_PASSWORD', 'admin'), engine_params={'echo': True} )() def check(session: Session, name=''): nonlocal critical for i in range(5): print("Lock await", name, i, datetime.now()) session.execute(text('SELECT pg_advisory_xact_lock(:lock)'), {'lock': 666}) print("Lock acquired", name, i, datetime.now()) self.assertFalse(critical) critical = True sleep(0.2) self.assertTrue(critical) critical = False session.commit() print("Lock released", name, i, datetime.now()) t1 = Thread(target=check, args=[create_session(), 't1']) t2 = Thread(target=check, args=[create_session(), 't2']) t3 = Thread(target=check, args=[create_session(), 't3']) t1.start() t2.start() t3.start() t1.join() t2.join() t3.join() self.assertTrue(True) def test_advisory_locked_transaction(self): uow_1 = UnitOfWork(**self._uows[1]) self._pg.add(uow_1) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name="test_content_name", created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() with self._repository.advisory_locked_transaction(): record = self._pg.query(DisassembleContentStatus).one() with self._repository.advisory_locked_transaction(): record = self._pg.query(DisassembleContentStatus).one() self.assertTrue(record) @parameterized.expand( [ (DateTimeRange(datetime(2019, 11, 24, 12), datetime(2019, 11, 24, 14)),), (None,), (0,) ] ) def test_get_record_to_replace(self, dt_range): uow_1 = UnitOfWork(**self._uows[1]) uow_2 = UnitOfWork(**self._uows[1]) uow_2.reprocess_id = 12345 if dt_range != 0: uow_1.timeslot = dt_range uow_2.timeslot = dt_range self._pg.add_all([uow_1, uow_2]) cs_1 = ContentStatus(**self._content_statuses[1]) uow_1.content_statuses.append(cs_1) cs_2 = ContentStatus(**self._content_statuses[1]) dcs_1 = DisassembleContentStatus( unit_of_work=uow_1, content_status=cs_1, content_name=cs_1.content_name, created_at=self._now, last_updated_at=self._now, disassemble_status=DisassembleStatusEnum.QUEUED ) self._pg.commit() record_to_replace = self._repository.get_record_to_replace( unit_of_work_id=uow_2.unit_of_work_id, context=cs_2.context, content_name=cs_2.content_name, ) assert record_to_replace == dcs_1