diff --git a/collector/collector/test/base.py b/collector/collector/test/base.py index d1d888e..6f0c4a8 100644 --- a/collector/collector/test/base.py +++ b/collector/collector/test/base.py @@ -68,20 +68,57 @@ class BaseTest(TestCase): class DbTest(BaseTest): + def get_migrations_dir(self): + return os.path.join(os.path.dirname(__file__), + '..', 'api', 'db', 'migrations') + def setUp(self): super(DbTest, self).setUp() - # Connection must be closed before DB migration - db.session.close() + # Cleaning all changes from the previous test + db.session.rollback() - # Cleaning DB. It useful in case of tests failure - directory = os.path.join(os.path.dirname(__file__), - '..', 'api', 'db', 'migrations') + directory = self.get_migrations_dir() with app.app_context(): try: flask_migrate.downgrade(directory=directory, revision='base') - except CommandError: - # Workaround for the first migration - pass + except CommandError as e: + app.logger.debug("DB migration downgrade failed: %s", e) + self.clean_db() flask_migrate.upgrade(directory=directory) + + def clean_db(self): + app.logger.debug("Cleaning DB without Alembic") + + # Removing tables + tables = db.session.execute( + "SELECT table_name FROM information_schema.tables " + "WHERE table_schema = 'public'") + table_names = list(item[0] for item in tables) + if table_names: + app.logger.debug("Removing tables: %s", table_names) + db.session.execute( + "DROP TABLE {0} CASCADE".format(','.join(table_names))) + + # Removing sequences + sequences = list(item[0] for item in db.session.execute( + "SELECT relname FROM pg_class WHERE relkind='S'")) + sequence_names = list(item[0] for item in sequences) + if sequence_names: + app.logger.debug("Removing sequences: %s", sequence_names) + db.session.execute( + "DROP SEQUENCE {0}".format(','.join(sequences))) + + # Removing enums + enums = db.session.execute( + "SELECT t.typname FROM pg_type t JOIN pg_catalog.pg_namespace n " + "ON n.oid = t.typnamespace WHERE n.nspname='public'") + enum_names = list(item[0] for item in enums) + if enum_names: + app.logger.debug("Removing types: %s", enum_names) + db.session.execute( + "DROP TYPE {0}".format(','.join(enum_names))) + + # Committing DDL changes + db.session.commit() diff --git a/collector/collector/test/db/test_migration.py b/collector/collector/test/db/test_migration.py new file mode 100644 index 0000000..b1e17d3 --- /dev/null +++ b/collector/collector/test/db/test_migration.py @@ -0,0 +1,43 @@ +# Copyright 2016 Mirantis, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +from alembic.util import CommandError +import flask_migrate + +from collector.test.base import DbTest + +from collector.api.app import app +from collector.api.app import db + + +class TestMigration(DbTest): + + def test_clean_db(self): + # Crashing alembic versions history + db.session.execute("UPDATE alembic_version SET version_num='x'") + db.session.commit() + + migrations_dir = self.get_migrations_dir() + with app.app_context(): + # Checking migrations are broken + self.assertRaises( + CommandError, flask_migrate.downgrade, + directory=migrations_dir, revision='base' + ) + + self.clean_db() + + # Checking migrations flow is fixed + flask_migrate.downgrade(directory=migrations_dir, revision='base') + flask_migrate.upgrade(directory=migrations_dir)