Support unique labels for alembic branches

We plan to have all projects that integrate with Neutron use the same
alembic environment (Neutron's) instead of having a separate alembic
env.py in each repo. This means that each project will need to label
their contract and expand branches uniquely.

With this patch we add support for branch labels of the form
  'networking-foo-contract'
  'networking-foo-expand'

We also support the old labels ('contract' and 'expand') with the
intent of backporting this patch to Mitaka so that projects can be
backwards compatible when they update their alembic branch labels. A
best effort will be made to backport this patch to Liberty also.

Support has been removed for label names starting with a "liberty_"
prefix. Such labels temporarily existed during the Liberty cycle but
have now all been cleaned up.

Partial-Bug: #1497830

Change-Id: Ibf38b3125d5d404d57f3e8ef0c7ecfa56eaf8042

doc/source/devref/alembic_migrations.rst

@@ -418,6 +418,20 @@ following directive should be added in the contraction script::
     depends_on = ('<expansion-revision>',)
 
 
+Expand and Contract Branch Labeling
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Before the Newton release, neutron and each sub-project had its own alembic
+environment and each repo could re-use the labels ``contract`` and ``expand``
+for their alembic branches. With the Newton release, all neutron sub-projects
+use neutron's alembic environment, and this requires globally unique branch
+labels.
+
+To be compatible with the Newton release of neutron, all projects must use
+unique alembic branch labels of the form ``<project>-contract`` and
+``<project>-expand``.
+
+
 HEAD files for conflict management
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

neutron/db/migration/alembic_migrations/versions/liberty/contract/30018084ec99_initial.py

@@ -25,7 +25,7 @@ from neutron.db.migration import cli
 # revision identifiers, used by Alembic.
 revision = '30018084ec99'
 down_revision = 'kilo'
-branch_labels = (cli.CONTRACT_BRANCH,)
+branch_labels = ('neutron-' + cli.CONTRACT_BRANCH,)
 
 
 def upgrade():

neutron/db/migration/alembic_migrations/versions/liberty/expand/354db87e3225_nsxv_vdr_metadata.py

@@ -30,7 +30,7 @@ from neutron.db.migration import cli
 # revision identifiers, used by Alembic.
 revision = '354db87e3225'
 down_revision = 'kilo'
-branch_labels = (cli.EXPAND_BRANCH,)
+branch_labels = ('neutron-' + cli.EXPAND_BRANCH,)
 
 
 def upgrade():

neutron/db/migration/autogen.py

@@ -54,7 +54,8 @@ def _migration_script_ops(context, directive, phase):
     version_path = cli._get_version_branch_path(
         context.config, release=cli.CURRENT_RELEASE, branch=phase)
     autogen_kwargs = {}
-    cli._check_bootstrap_new_branch(phase, version_path, autogen_kwargs)
+    cli._check_bootstrap_new_branch(context.config, phase, version_path,
+                                    autogen_kwargs)
 
     op = ops.MigrationScript(
         new_rev_id(),

neutron/db/migration/cli.py

@@ -93,11 +93,14 @@ CONF.register_cli_opts(_db_opts, 'database')
 
 
 def do_alembic_command(config, cmd, revision=None, desc=None, **kwargs):
+    project = config.get_main_option('neutron_project')
     args = []
     if revision:
+        # We use unique branch labels from Newton onwards.
+        if revision.split('@')[0] in MIGRATION_BRANCHES:
+            revision = '-'.join([project, revision])
         args.append(revision)
 
-    project = config.get_main_option('neutron_project')
     if desc:
         alembic_util.msg(_('Running %(cmd)s (%(desc)s) for %(project)s ...') %
                          {'cmd': cmd, 'desc': desc, 'project': project})
@@ -142,13 +145,13 @@ def add_branch_options(parser):
 def _find_milestone_revisions(config, milestone, branch=None):
     """Return the revision(s) for a given milestone."""
     script = alembic_script.ScriptDirectory.from_config(config)
-    return [
-        (m.revision, label)
-        for m in _get_revisions(script)
-        for label in (m.branch_labels or [None])
-        if milestone in getattr(m.module, 'neutron_milestone', []) and
-        (branch is None or branch in m.branch_labels)
-    ]
+    milestone_revisions = []
+    for m in _get_revisions(script):
+        for branch_label in (m.branch_labels or [None]):
+            if milestone in getattr(m.module, 'neutron_milestone', []):
+                if branch is None or branch_label.endswith(branch):
+                    milestone_revisions.append((m.revision, branch))
+    return milestone_revisions
 
 
 def do_upgrade(config, cmd):
@@ -161,11 +164,11 @@ def do_upgrade(config, cmd):
 
     if CONF.command.expand:
         branch = EXPAND_BRANCH
-        revision = _get_branch_head(EXPAND_BRANCH)
+        revision = _get_branch_head(config, EXPAND_BRANCH)
 
     elif CONF.command.contract:
         branch = CONTRACT_BRANCH
-        revision = _get_branch_head(CONTRACT_BRANCH)
+        revision = _get_branch_head(config, CONTRACT_BRANCH)
 
     elif not CONF.command.revision and not CONF.command.delta:
         raise SystemExit(_('You must provide a revision or relative delta'))
@@ -217,14 +220,26 @@ def do_stamp(config, cmd):
                        sql=CONF.command.sql)
 
 
-def _get_branch_head(branch):
+def _get_branch_head(config, branch):
     '''Get the latest @head specification for a branch.'''
+    script_dir = alembic_script.ScriptDirectory.from_config(config)
+    revs = script_dir.revision_map.get_revisions('heads')
+    for rev in revs:
+        for branch_label in rev.branch_labels:
+            # For forwards and backwards compatibility we handle branch
+            # names of either
+            #    'contract/expand'
+            # or
+            #    'subproject-contract/subproject-expand'.
+            if branch_label.endswith(branch):
+                branch = branch_label
+                break
     return '%s@head' % branch
 
 
-def _check_bootstrap_new_branch(branch, version_path, addn_kwargs):
+def _check_bootstrap_new_branch(config, branch, version_path, addn_kwargs):
     addn_kwargs['version_path'] = version_path
-    addn_kwargs['head'] = _get_branch_head(branch)
+    addn_kwargs['head'] = _get_branch_head(config, branch)
     if not os.path.exists(version_path):
         # Bootstrap initial directory structure
         utils.ensure_dir(version_path)
@@ -251,7 +266,7 @@ def do_revision(config, cmd):
             args = copy.copy(kwargs)
             version_path = _get_version_branch_path(
                 config, release=CURRENT_RELEASE, branch=branch)
-            _check_bootstrap_new_branch(branch, version_path, args)
+            _check_bootstrap_new_branch(config, branch, version_path, args)
             do_alembic_command(config, cmd, **args)
     else:
         # autogeneration code will take care of enforcing proper directories
@@ -260,32 +275,15 @@ def do_revision(config, cmd):
     update_head_files(config)
 
 
-def _get_release_labels(labels):
-    result = set()
-    for label in labels:
-        # release labels were introduced Liberty for a short time and dropped
-        # in that same release cycle
-        result.add('%s_%s' % (migration.LIBERTY, label))
-    return result
-
-
 def _compare_labels(revision, expected_labels):
     # validate that the script has expected labels only
-    bad_labels = revision.branch_labels - expected_labels
+    expected_branch_labels = set()
+    for label in revision.branch_labels:
+        for expected_label in expected_labels:
+            if label.endswith(expected_label):
+                expected_branch_labels.add(label)
+    bad_labels = revision.branch_labels - expected_branch_labels
     if bad_labels:
-        # NOTE(ihrachyshka): this hack is temporary to accommodate those
-        # projects that already initialized their branches with liberty_*
-        # labels. Let's notify them about the deprecation for now and drop it
-        # later.
-        bad_labels_with_release = (revision.branch_labels -
-                                   _get_release_labels(expected_labels))
-        if not bad_labels_with_release:
-            alembic_util.warn(
-                _('Release aware branch labels (%s) are deprecated. '
-                  'Please switch to expand@ and contract@ '
-                  'labels.') % bad_labels)
-            return
-
         script_name = os.path.basename(revision.path)
         alembic_util.err(
             _('Unexpected label for script %(script_name)s: %(labels)s') %
@@ -339,6 +337,14 @@ def _get_revisions(script):
     return list(script.walk_revisions(base='base', head='heads'))
 
 
+def _get_branch_type(revision):
+    for branch_label in revision.branch_labels:
+        if branch_label.endswith(CONTRACT_BRANCH):
+            return CONTRACT_BRANCH
+        if branch_label.endswith(EXPAND_BRANCH):
+            return EXPAND_BRANCH
+
+
 def _get_branch_points(script):
     branchpoints = []
     for revision in _get_revisions(script):
@@ -352,10 +358,7 @@ def _get_heads_map(config):
     heads = script.get_heads()
     head_map = {}
     for head in heads:
-        if CONTRACT_BRANCH in script.get_revision(head).branch_labels:
-            head_map[CONTRACT_BRANCH] = head
-        else:
-            head_map[EXPAND_BRANCH] = head
+        head_map[_get_branch_type(script.get_revision(head))] = head
     return head_map
 
 

neutron/tests/functional/db/test_migrations.py

@@ -237,7 +237,12 @@ class TestModelsMigrationsMysql(_TestModelsMigrations,
             script = alembic_script.ScriptDirectory.from_config(
                 self.alembic_config)
             for m in list(script.walk_revisions(base='base', head='heads')):
-                branches = m.branch_labels or []
+                branches = set()
+                for bl in m.branch_labels:
+                    if bl.endswith(migration.CONTRACT_BRANCH):
+                        branches.add(migration.CONTRACT_BRANCH)
+                    elif bl.endswith(migration.EXPAND_BRANCH):
+                        branches.add(migration.EXPAND_BRANCH)
                 if migration.CONTRACT_BRANCH in branches:
                     method_name = 'contract_creation_exceptions'
                     exceptions_dict = creation_exceptions

neutron/tests/unit/db/test_migration.py

@@ -148,7 +148,28 @@ class TestCli(base.BaseTestCase):
                                                   attrs=attrs)
             cli.migration_entrypoints[project] = entrypoint
 
-    def _main_test_helper(self, argv, func_name, exp_kwargs=[{}]):
+        self.old_labels_revs = {
+            'a': FakeRevision(labels={cli.CONTRACT_BRANCH}),
+            'b': FakeRevision(labels={cli.EXPAND_BRANCH})
+        }
+        self.new_labels_revs = {
+            'x': FakeRevision(labels={'foo-' + cli.CONTRACT_BRANCH}),
+            'y': FakeRevision(labels={'foo-' + cli.EXPAND_BRANCH})
+        }
+
+    @staticmethod
+    def _old_and_new_label_combination_helper(test_method):
+        old_cb = cli.CONTRACT_BRANCH
+        new_cb = 'foo-' + cli.CONTRACT_BRANCH
+        old_eb = cli.EXPAND_BRANCH
+        new_eb = 'foo-' + cli.EXPAND_BRANCH
+        # Test all combinations of old and new branch labels.
+        test_method(old_cb, old_eb)
+        test_method(old_cb, new_eb)
+        test_method(new_cb, old_eb)
+        test_method(new_cb, new_eb)
+
+    def _main_test_helper(self, argv, func_name, exp_kwargs=None):
         with mock.patch.object(sys, 'argv', argv),\
                 mock.patch.object(cli, 'run_sanity_checks'),\
                 mock.patch.object(cli, 'validate_revisions'):
@@ -163,7 +184,7 @@ class TestCli(base.BaseTestCase):
 
             self.do_alembic_cmd.assert_has_calls(
                 [mock.call(mock.ANY, func_name, **_append_version_path(kwargs))
-                 for kwargs in exp_kwargs]
+                 for kwargs in (exp_kwargs or [{}])]
             )
 
     def test_stamp(self):
@@ -207,10 +228,11 @@ class TestCli(base.BaseTestCase):
             self._main_test_helper(['prog', 'check_migration'], 'branches')
             self.assertEqual(len(self.projects), validate.call_count)
 
-    def _test_database_sync_revision(self, separate_branches=True):
-        with mock.patch.object(cli, 'update_head_files') as update:
-            if separate_branches:
-                mock.patch('os.path.exists').start()
+    def _test_database_sync_revision(self, revs):
+        with mock.patch.object(cli, 'update_head_files') as update,\
+                mock.patch('alembic.script.ScriptDirectory.from_config') as fc:
+            fc.return_value.get_revisions.side_effect = revs
+            mock.patch('os.path.exists').start()
             expected_kwargs = [{
                 'message': 'message', 'sql': False, 'autogenerate': True,
             }]
@@ -226,7 +248,7 @@ class TestCli(base.BaseTestCase):
                 'message': 'message',
                 'sql': True,
                 'autogenerate': False,
-                'head': cli._get_branch_head(branch)
+                'head': cli._get_branch_head(self.configs[0], branch)
             } for branch in cli.MIGRATION_BRANCHES]
             for kwarg in expected_kwargs:
                 kwarg['autogenerate'] = False
@@ -264,12 +286,11 @@ class TestCli(base.BaseTestCase):
             )
             self.assertEqual(len(self.projects), update.call_count)
 
-    def test_database_sync_revision(self):
-        self._test_database_sync_revision()
+    def test_database_sync_revision_old_branch_labels(self):
+        self._test_database_sync_revision(self.old_labels_revs)
 
-    def test_database_sync_revision_no_branches(self):
-        # Test that old branchless approach is still supported
-        self._test_database_sync_revision(separate_branches=False)
+    def test_database_sync_revision_new_branch_labels(self):
+        self._test_database_sync_revision(self.new_labels_revs)
 
     def test_upgrade_revision(self):
         self._main_test_helper(
@@ -292,14 +313,22 @@ class TestCli(base.BaseTestCase):
             [{'desc': None, 'revision': 'kilo+3', 'sql': False}]
         )
 
-    def test_upgrade_expand(self):
-        self._main_test_helper(
-            ['prog', 'upgrade', '--expand'],
-            'upgrade',
-            [{'desc': cli.EXPAND_BRANCH,
-              'revision': 'expand@head',
-              'sql': False}]
-        )
+    def _test_upgrade_expand(self, revs):
+        with mock.patch('alembic.script.ScriptDirectory.from_config') as fc:
+            fc.return_value.get_revisions.side_effect = revs
+            self._main_test_helper(
+                ['prog', 'upgrade', '--expand'],
+                'upgrade',
+                [{'desc': cli.EXPAND_BRANCH,
+                  'revision': 'expand@head',
+                  'sql': False}]
+            )
+
+    def test_upgrade_expand_old_labels(self):
+        self._test_upgrade_expand(self.old_labels_revs)
+
+    def test_upgrade_expand_new_labels(self):
+        self._test_upgrade_expand(self.new_labels_revs)
 
     def test_upgrade_expand_contract_are_mutually_exclusive(self):
         with testlib_api.ExpectedException(SystemExit):
@@ -342,20 +371,31 @@ class TestCli(base.BaseTestCase):
     def test_upgrade_contract_conflicts_with_delta(self):
         self._test_upgrade_conflicts_with_delta('contract')
 
-    def test_upgrade_contract(self):
-        self._main_test_helper(
-            ['prog', 'upgrade', '--contract'],
-            'upgrade',
-            [{'desc': cli.CONTRACT_BRANCH,
-              'revision': 'contract@head',
-              'sql': False}]
-        )
+    def _test_upgrade_contract(self, revs):
+        with mock.patch('alembic.script.ScriptDirectory.from_config') as fc:
+            fc.return_value.get_revisions.side_effect = revs
+            self._main_test_helper(
+                ['prog', 'upgrade', '--contract'],
+                'upgrade',
+                [{'desc': cli.CONTRACT_BRANCH,
+                  'revision': 'contract@head',
+                  'sql': False}]
+            )
+
+    def test_upgrade_contract_old_labels(self):
+        self._test_upgrade_contract(self.old_labels_revs)
+
+    def test_upgrade_contract_new_labels(self):
+        self._test_upgrade_contract(self.new_labels_revs)
 
     @mock.patch('alembic.script.ScriptDirectory.walk_revisions')
-    def test_upgrade_milestone_expand_before_contract(self, walk_mock):
-        c_revs = [FakeRevision(labels={cli.CONTRACT_BRANCH}) for r in range(5)]
+    def _test_upgrade_milestone_expand_before_contract(self,
+                                                       contract_label,
+                                                       expand_label,
+                                                       walk_mock):
+        c_revs = [FakeRevision(labels={contract_label}) for r in range(5)]
         c_revs[1].module.neutron_milestone = [migration.LIBERTY]
-        e_revs = [FakeRevision(labels={cli.EXPAND_BRANCH}) for r in range(5)]
+        e_revs = [FakeRevision(labels={expand_label}) for r in range(5)]
         e_revs[3].module.neutron_milestone = [migration.LIBERTY]
         walk_mock.return_value = c_revs + e_revs
         self._main_test_helper(
@@ -369,6 +409,10 @@ class TestCli(base.BaseTestCase):
               'sql': False}]
         )
 
+    def test_upgrade_milestone_expand_before_contract(self):
+        self._old_and_new_label_combination_helper(
+            self._test_upgrade_milestone_expand_before_contract)
+
     def assert_command_fails(self, command):
         # Avoid cluttering stdout with argparse error messages
         mock.patch('argparse.ArgumentParser._print_message').start()
@@ -388,8 +432,8 @@ class TestCli(base.BaseTestCase):
     def test_upgrade_rejects_delta_with_relative_revision(self):
         self.assert_command_fails(['prog', 'upgrade', '+2', '--delta', '3'])
 
-    def _test_validate_head_files_helper(self, heads, contract_head='',
-                                         expand_head=''):
+    def _test_validate_head_files_common(self, heads, revs,
+                                         contract_head='', expand_head=''):
         fake_config = self.configs[0]
         head_files_not_exist = (contract_head == expand_head == '')
         with mock.patch('alembic.script.ScriptDirectory.from_config') as fc,\
@@ -400,9 +444,6 @@ class TestCli(base.BaseTestCase):
                 os_mock.return_value = True
 
             fc.return_value.get_heads.return_value = heads
-
-            revs = {heads[0]: FakeRevision(labels=cli.CONTRACT_BRANCH),
-                    heads[1]: FakeRevision(labels=cli.EXPAND_BRANCH)}
             fc.return_value.get_revision.side_effect = revs.__getitem__
             mock_open_con = self.useFixture(
                 tools.OpenFixture(cli._get_contract_head_file_path(
@@ -433,6 +474,23 @@ class TestCli(base.BaseTestCase):
             if not head_files_not_exist:
                 fc.assert_called_once_with(fake_config)
 
+    def _test_validate_head_files_helper(self, heads,
+                                         contract_head='', expand_head=''):
+        old_labels_revs = {
+            heads[0]: FakeRevision(labels={cli.CONTRACT_BRANCH}),
+            heads[1]: FakeRevision(labels={cli.EXPAND_BRANCH})
+        }
+        new_labels_revs = {
+            heads[0]: FakeRevision(labels={'foo-' + cli.CONTRACT_BRANCH}),
+            heads[1]: FakeRevision(labels={'foo-' + cli.EXPAND_BRANCH})
+        }
+        self._test_validate_head_files_common(heads, old_labels_revs,
+                                              contract_head=contract_head,
+                                              expand_head=expand_head)
+        self._test_validate_head_files_common(heads, new_labels_revs,
+                                              contract_head=contract_head,
+                                              expand_head=expand_head)
+
     def test_validate_head_files_success(self):
         self._test_validate_head_files_helper(['a', 'b'], contract_head='a',
                                               expand_head='b')
@@ -445,8 +503,14 @@ class TestCli(base.BaseTestCase):
                                               expand_head='d')
 
     @mock.patch.object(fileutils, 'delete_if_exists')
-    def test_update_head_files_success(self, *mocks):
-        heads = ['a', 'b']
+    def _test_update_head_files_success(self, revs, *mocks):
+        contract_head = expand_head = None
+        for rev, revision in revs.items():
+            for branch_label in revision.branch_labels:
+                if branch_label.endswith(cli.CONTRACT_BRANCH):
+                    contract_head = rev
+                if branch_label.endswith(cli.EXPAND_BRANCH):
+                    expand_head = rev
         mock_open_con = self.useFixture(
                     tools.OpenFixture(cli._get_contract_head_file_path(
                         self.configs[0]))).mock_open
@@ -454,14 +518,13 @@ class TestCli(base.BaseTestCase):
             tools.OpenFixture(cli._get_expand_head_file_path(
                 self.configs[0]))).mock_open
         with mock.patch('alembic.script.ScriptDirectory.from_config') as fc:
-            fc.return_value.get_heads.return_value = heads
-            revs = {heads[0]: FakeRevision(labels=cli.CONTRACT_BRANCH),
-                    heads[1]: FakeRevision(labels=cli.EXPAND_BRANCH)}
+            fc.return_value.get_heads.return_value = list(revs)
             fc.return_value.get_revision.side_effect = revs.__getitem__
             cli.update_head_files(self.configs[0])
             mock_open_con.return_value.write.assert_called_with(
-                heads[0] + '\n')
-            mock_open_ex.return_value.write.assert_called_with(heads[1] + '\n')
+                contract_head + '\n')
+            mock_open_ex.return_value.write.assert_called_with(
+                expand_head + '\n')
 
             old_head_file = cli._get_head_file_path(
                 self.configs[0])
@@ -473,6 +536,12 @@ class TestCli(base.BaseTestCase):
             self.assertIn(mock.call(old_heads_file),
                           delete_if_exists.call_args_list)
 
+    def test_update_head_files_success_old_labels(self):
+        self._test_update_head_files_success(self.old_labels_revs)
+
+    def test_update_head_files_success_new_labels(self):
+        self._test_update_head_files_success(self.new_labels_revs)
+
     def test_get_project_base(self):
         config = alembic_config.Config()
         config.set_main_option('script_location', 'a.b.c:d')
@@ -589,7 +658,7 @@ class TestCli(base.BaseTestCase):
 
     @mock.patch('alembic.script.ScriptDirectory.walk_revisions')
     def test__get_branch_points(self, walk_mock):
-        revisions = [FakeRevision(is_branch_point=tools.get_random_boolean)
+        revisions = [FakeRevision(is_branch_point=tools.get_random_boolean())
                      for i in range(50)]
         walk_mock.return_value = revisions
         script_dir = alembic_script.ScriptDirectory.from_config(
@@ -663,8 +732,10 @@ class TestCli(base.BaseTestCase):
         )
 
         directives = [migration_script]
+        context = mock.Mock()
+        context.config = self.configs[0]
         autogen.process_revision_directives(
-            mock.Mock(), mock.Mock(), directives
+            context, mock.Mock(), directives
         )
 
         expand = directives[0]
@@ -699,23 +770,33 @@ class TestCli(base.BaseTestCase):
         )
 
     @mock.patch('alembic.script.ScriptDirectory.walk_revisions')
-    def test__find_milestone_revisions_one_branch(self, walk_mock):
-        c_revs = [FakeRevision(labels={cli.CONTRACT_BRANCH}) for r in range(5)]
+    def _test__find_milestone_revisions_one_branch(self,
+                                                   contract_label,
+                                                   expand_label,
+                                                   walk_mock):
+        c_revs = [FakeRevision(labels={contract_label}) for r in range(5)]
         c_revs[1].module.neutron_milestone = [migration.LIBERTY]
 
         walk_mock.return_value = c_revs
         m = cli._find_milestone_revisions(self.configs[0], 'liberty',
-                                          cli.CONTRACT_BRANCH)
+                                          contract_label)
         self.assertEqual(1, len(m))
         m = cli._find_milestone_revisions(self.configs[0], 'liberty',
-                                          cli.EXPAND_BRANCH)
+                                          expand_label)
         self.assertEqual(0, len(m))
 
+    def test__find_milestone_revisions_one_branch(self):
+        self._old_and_new_label_combination_helper(
+            self._test__find_milestone_revisions_one_branch)
+
     @mock.patch('alembic.script.ScriptDirectory.walk_revisions')
-    def test__find_milestone_revisions_two_branches(self, walk_mock):
-        c_revs = [FakeRevision(labels={cli.CONTRACT_BRANCH}) for r in range(5)]
+    def _test__find_milestone_revisions_two_branches(self,
+                                                     contract_label,
+                                                     expand_label,
+                                                     walk_mock):
+        c_revs = [FakeRevision(labels={contract_label}) for r in range(5)]
         c_revs[1].module.neutron_milestone = [migration.LIBERTY]
-        e_revs = [FakeRevision(labels={cli.EXPAND_BRANCH}) for r in range(5)]
+        e_revs = [FakeRevision(labels={expand_label}) for r in range(5)]
         e_revs[3].module.neutron_milestone = [migration.LIBERTY]
 
         walk_mock.return_value = c_revs + e_revs
@@ -725,6 +806,10 @@ class TestCli(base.BaseTestCase):
         m = cli._find_milestone_revisions(self.configs[0], 'mitaka')
         self.assertEqual(0, len(m))
 
+    def test__find_milestone_revisions_two_branches(self):
+        self._old_and_new_label_combination_helper(
+            self._test__find_milestone_revisions_two_branches)
+
     @mock.patch('alembic.script.ScriptDirectory.walk_revisions')
     def test__find_milestone_revisions_branchless(self, walk_mock):
         revisions = [FakeRevision() for r in range(5)]