Source code for aiida.backends.sqlalchemy.testbase

# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
# pylint: disable=import-error,no-name-in-module
""" This module contains the codebase for the setUpClass and tearDown methods used
internally by the AiidaTestCase. This inherits only from 'object' to avoid
that it is picked up by the automatic discovery of tests
(It shouldn't, as it risks to destroy the DB if there are not the checks
in place, and these are implemented in the AiidaTestCase. """

from aiida.backends.testimplbase import AiidaTestImplementation


class SqlAlchemyTests(AiidaTestImplementation):
    """Base class to test SQLA-related functionalities."""
    connection = None

    def clean_db(self):
        from sqlalchemy.sql import table
        # pylint: disable=invalid-name
        DbGroupNodes = table('db_dbgroup_dbnodes')
        DbGroup = table('db_dbgroup')
        DbLink = table('db_dblink')
        DbNode = table('db_dbnode')
        DbLog = table('db_dblog')
        DbAuthInfo = table('db_dbauthinfo')
        DbUser = table('db_dbuser')
        DbComputer = table('db_dbcomputer')

        with self.backend.transaction() as session:
            session.execute(DbGroupNodes.delete())
            session.execute(DbGroup.delete())
            session.execute(DbLog.delete())
            session.execute(DbLink.delete())
            session.execute(DbNode.delete())
            session.execute(DbAuthInfo.delete())
            session.execute(DbComputer.delete())
            session.execute(DbUser.delete())
            session.commit()