class TestDBLogger(unittest.TestCase):
def setUp(self):
self.cm = ConnectionManager(default_db)
self.cm.connect()
def tearDown(self):
self.cm.close()
def test_create_and_finalize(self):
db_logger = DatabaseLogger(cm=self.cm, report_name="TestDBLogger")
db_logger.create_record()
self.assertIsInstance(db_logger.report_log.primary_key, int)
self.assertTrue(db_logger.report_log.inserted)
self.assertFalse(db_logger.report_log.updated)
db_logger.finalize_record()
self.assertTrue(db_logger.report_log.updated)
def test_log_recipients(self):
test_recipients = ["*****@*****.**", "*****@*****.**"]
db_logger = DatabaseLogger(cm=self.cm, report_name="TestDBLogger")
db_logger.create_record()
db_logger.finalize_record(test_recipients)
sql = f"select recipients from {self.cm.schema}.report_logs where id = {db_logger.report_log.primary_key}"
with QueryExecutor(default_db) as qe:
result = qe.execute_query(sql=sql)
self.assertEqual(result.result_data[0]['recipients'], "; ".join(test_recipients))
class TestReportWriter(unittest.TestCase):
def setUp(self):
self.cm = ConnectionManager(default_db)
self.cm.connect()
def tearDown(self):
self.cm.close()
def test_create_workbook(self):
writer = ReportWriter("Test Report")
writer.build_file()
writer.close_workbook()
self.assertIsNotNone(writer.workbook_builder)
self.assertTrue(os.path.isfile(writer.report_file))
os.unlink(writer.report_file)
def test_create_worksheet(self):
writer = ReportWriter("Test Report 2")
writer.build_file()
self.assertEqual(writer.record_count, 0)
test_query = TEST_QUERY.format(self.cm.schema)
writer.create_worksheet_from_query(self.cm, "sheet1", sql=test_query)
writer.close_workbook()
self.assertTrue(writer.record_count > 0)
self.assertTrue(writer.logger.error_buffer.empty)
def test_invalid_sheet_name_raises_error(self):
"""Should log an error if attempt to add worksheet with invalid name"""
writer = ReportWriter("Test Report 3")
writer.build_file()
test_query = TEST_QUERY.format(self.cm.schema)
writer.create_worksheet_from_query(self.cm, "sheet/1", sql=test_query)
writer.close_workbook()
self.assertFalse(writer.logger.error_buffer.empty)
def test_close_connection_successfully(self):
# TODO This test only works if the 'Test' db is defined in config.ini.
# TODO Should instead create a test db fixture of some kind.
cm = ConnectionManager(db='Test')
cm.connect()
cm.conn.close()
with self.assertRaisesRegex(StatementError, "ResourceClosedError"):
cm.conn.execute("select * from test;")
class TestReportActiveChecker(unittest.TestCase):
def setUp(self):
self.cm = ConnectionManager(default_db)
self.cm.connect()
def tearDown(self):
self.cm.close()
def test_active(self):
active_checker = ReportActiveChecker(self.cm, 'test_report_active')
inactive_checker = ReportActiveChecker(self.cm, 'test_report_inactive')
self.assertTrue(active_checker)
self.assertFalse(inactive_checker)
class TestRecipientsChecker(unittest.TestCase):
def setUp(self):
self.cm = ConnectionManager(default_db)
self.cm.connect()
def tearDown(self):
self.cm.close()
def test_get_recipients(self):
checker = RecipientsChecker(cm=self.cm, report_name='test_report_active')
checker.assertFalse(checker.to_recipients)
checker.assertFalse(checker.cc_recipients)
checker.get_recipients()
checker.assertTrue(checker.to_recipients)
checker.assertTrue(checker.cc_recipients)
def execute_statement(self, statement: sa.sql.base.Executable):
with ConnectionManager(self.database) as cm:
proxy = cm.conn.execute(statement)
if proxy.cursor:
result = proxy.fetchall()
else:
result = None
return result
def teardown_test_db():
db = config['Default']['database']
try:
os.unlink('test.db')
except FileNotFoundError:
pass
if config[db]['rdbms'] == 'mysql':
with ConnectionManager(db) as cm:
metadata.drop_all(cm.engine)
test_metadata.drop_all(cm.engine)
def send_unschedulable_alert(report):
with ConnectionManager('Celery-Docker') as cm:
checker = RecipientsChecker(cm, 'alert')
checker.get_recipients()
alert = Alert(
subject=f"Unschedulable Task: {report['report_name']}",
message=
"Task is active but cannot be scheduled. You must provide all cron schedule attributes, or none.",
recipients=checker.to_recipients)
alert.send()
def test_allow_valid_rdbms_types(self):
"""Raise an error if an unsupported RDBMS is specified."""
cm = ConnectionManager()
cm.db = 'Fake_DB'
cm.config.add_section('Fake_DB')
cm.config.set('Fake_DB', 'rdbms', 'fakedb')
cm.unpack_params()
with self.assertRaisesRegex(ValueError, "Unsupported RDBMS"):
cm.connect()
def setup_test_db():
db = config['Default']['database']
if config[db]['rdbms'] == 'sqlite':
cm = ConnectionManager(db)
cm.connect()
metadata.create_all(cm.engine)
test_metadata.create_all(cm.engine)
create_fixtures(cm)
cm.close()
if config[db]['rdbms'] == 'mysql':
with ConnectionManager(db) as cm:
metadata.create_all(cm.engine)
test_metadata.create_all(cm.engine)
create_fixtures(cm)
def test_raise_error_if_no_db(self):
"""Raise an error if no db is provided"""
cm = ConnectionManager()
with self.assertRaisesRegex(ValueError, "db attribute not set"):
cm.connect()
def test_context_manager(self):
with ConnectionManager(db='Test') as cm:
self.assertFalse(cm.closed())
cm.conn.execute("select * from sys.flarp;")
self.assertTrue(cm.closed())
def setUp(self):
self.cm = ConnectionManager(default_db)
self.cm.connect()
def setUp(self):
db = config['Default']['database']
cm = ConnectionManager(db)
cm.connect()
self.cm = cm