def setUp(self):
"""Go to a specific schema version before running tests."""
from aiida.backends.djsite import get_scoped_session
from aiida.orm import autogroup
self.current_autogroup = autogroup.CURRENT_AUTOGROUP
autogroup.CURRENT_AUTOGROUP = None
assert self.migrate_from and self.migrate_to, \
"TestCase '{}' must define migrate_from and migrate_to properties".format(type(self).__name__)
self.migrate_from = [(self.app, self.migrate_from)]
self.migrate_to = [(self.app, self.migrate_to)]
executor = MigrationExecutor(connection)
self.apps = executor.loader.project_state(self.migrate_from).apps
self.schema_editor = connection.schema_editor()
# Before running the migration, make sure we close the querybuilder session which may still contain references
# to objects whose mapping may be invalidated after resetting the schema to an older version. This can block
# the migrations so we first expunge those objects by closing the session.
get_scoped_session().close()
# Reverse to the original migration
with Capturing():
executor.migrate(self.migrate_from)
self.DbLink = self.apps.get_model('db', 'DbLink')
self.DbNode = self.apps.get_model('db', 'DbNode')
self.DbUser = self.apps.get_model('db', 'DbUser')
self.DbUser.objects.all().delete()
self.default_user = self.DbUser(1, 'aiida@localhost')
self.default_user.save()
try:
self.setUpBeforeMigration()
# Run the migration to test
executor = MigrationExecutor(connection)
executor.loader.build_graph()
with Capturing():
executor.migrate(self.migrate_to)
self.apps = executor.loader.project_state(self.migrate_to).apps
except Exception:
# Bring back the DB to the correct state if this setup part fails
import traceback
traceback.print_stack()
self._revert_database_schema()
raise
def get_session():
"""Return a database session that can be used by the `QueryBuilder` to perform its query.
If there is an exception within the context then the changes will be rolled back and the state will
be as before entering. Transactions can be nested.
:return: an instance of :class:`sqlalchemy.orm.session.Session`
"""
from aiida.backends.djsite import get_scoped_session
return get_scoped_session()