def setUp(self):
"""Setups the global database environment."""
# Set up database before calling super.setUp(), in case any other mixin
# depends on db during its setup.
db_obj, cleanup = self.CreateDatabase()
patcher = mock.patch.object(data_store, "REL_DB",
db.DatabaseValidationWrapper(db_obj))
patcher.start()
self.addCleanup(patcher.stop)
if cleanup:
self.addCleanup(cleanup)
# In case a test registers a message handler, unregister it.
self.addCleanup(data_store.REL_DB.UnregisterMessageHandler)
super().setUp()
def Wrapper(*args, **kwargs):
db = abstract_db.DatabaseValidationWrapper(mem.InMemoryDB())
with mock.patch.object(data_store, "REL_DB", db):
func(*(args + (db,)), **kwargs)