diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index f68ed1e4d2181..2e78351fc4cfc 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -20,7 +20,9 @@ import unittest from alembic.autogenerate import compare_metadata +from alembic.config import Config from alembic.migration import MigrationContext +from alembic.script import ScriptDirectory from sqlalchemy import MetaData from airflow.models import Base as airflow_base @@ -97,3 +99,12 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self): diff, 'Database schema and SQLAlchemy model are not in sync: ' + str(diff) ) + + def test_only_single_head_revision_in_migrations(self): + config = Config() + config.set_main_option("script_location", "airflow:migrations") + script = ScriptDirectory.from_config(config) + + # This will raise if there are multiple heads + # To resolve, use the command `alembic merge` + script.get_current_head()