updated generate_deployments script to use new InstanceConfig

paasta_tools/generate_deployments_for_service.py

@@ -48,9 +48,11 @@
 
 from paasta_tools import remote_git
 from paasta_tools.utils import atomic_file_write
-from paasta_tools.utils import get_paasta_branch
 from paasta_tools.utils import get_git_url
-
+from paasta_tools.utils import list_clusters
+from paasta_tools.utils import get_service_instance_list
+from paasta_tools.marathon_tools import load_marathon_service_config
+from paasta_tools.chronos_tools import load_chronos_job_config
 
 log = logging.getLogger('__main__')
 logging.basicConfig()
@@ -70,35 +72,33 @@ def parse_args():
     return args
 
 
-def get_branches_from_config_file(file_dir, filename):
-    """Get all branches defined in a single service configuration file.
-    A branch is defined for an instance if it has a 'branch' key, or
-    the branch name is paasta-{cluster}.{instance},
-    where cluster is the cluster the marathon or chronos file is defined for
-    (i.e. marathon-hab.yaml is for hab), and instance is the
-    instance name.
-
-    :param file_dir: The directory that the filename argument is in
-    :param filename: The name of the service configuration file to read from
-    :returns: A set of branch names listed in the configuration file
-    """
-    valid_branches = set([])
-    config = service_configuration_lib.read_service_information(os.path.join(file_dir, filename))
-    for instance in config:
-        target_branch = None
-        if 'branch' in config[instance]:
-            target_branch = config[instance]['branch']
-        else:
-            try:
-                # cluster may contain dashes (and frequently does) so
-                # reassemble the cluster after stripping the chronos/marathon prefix
-                cluster = '-'.join(filename.split('-')[1:]).split('.')[0]
-                target_branch = get_paasta_branch(cluster, instance)
-            except IndexError:
-                pass
-        if target_branch:
-            valid_branches.add(target_branch)
-    return valid_branches
+def get_instance_config_for_service(soa_dir, service):
+    for cluster in list_clusters(
+        service=service,
+        soa_dir=soa_dir,
+    ):
+        for _, instance in get_service_instance_list(
+            service=service,
+            cluster=cluster,
+            instance_type='marathon',
+        ):
+            yield load_marathon_service_config(
+                service=service,
+                instance=instance,
+                cluster=cluster,
+                soa_dir=soa_dir,
+            )
+        for _, instance in get_service_instance_list(
+            service=service,
+            cluster=cluster,
+            instance_type='chronos',
+        ):
+            yield load_chronos_job_config(
+                service=service,
+                instance=instance,
+                cluster=cluster,
+                soa_dir=soa_dir,
+            )
 
 
 def get_branches_for_service(soa_dir, service):
@@ -108,11 +108,11 @@ def get_branches_for_service(soa_dir, service):
     :param service: The service name to get branches for
     :returns: A list of branches defined in instances for the service
     """
-    valid_branches = set([])
-    working_dir = os.path.join(soa_dir, service)
-    for fname in os.listdir(working_dir):
-        if fname.startswith('marathon-') or fname.startswith('chronos-'):
-            valid_branches = valid_branches.union(get_branches_from_config_file(working_dir, fname))
+    valid_branches = set([config.get_deploy_group() for config in get_instance_config_for_service(
+        soa_dir=soa_dir,
+        service=service,
+    )])
+
     return valid_branches
 
 

tests/test_generate_deployments_for_service.py

@@ -15,51 +15,51 @@
 import generate_deployments_for_service
 import mock
 import contextlib
-
-
-def test_get_branches_from_config_file():
-    fake_dir = '/nail/etc/soa/dir'
-    fake_fname = 'marathon-boston-hab-with-extra-dashes.yaml'
-    fake_config = {'test1': {'branch': 'ranch'}, 'test2': {}}
-    expected = set(['ranch', 'paasta-boston-hab-with-extra-dashes.test2'])
-    with contextlib.nested(
-        mock.patch('service_configuration_lib.read_service_information',
-                   return_value=fake_config),
-        mock.patch('os.path.join', return_value='internal_operations')
-    ) as (
-        read_info_patch,
-        join_patch
-    ):
-        actual = generate_deployments_for_service.get_branches_from_config_file(fake_dir, fake_fname)
-        assert expected == actual
-        join_patch.assert_called_once_with(fake_dir, fake_fname)
-        read_info_patch.assert_called_once_with('internal_operations')
+from paasta_tools.marathon_tools import MarathonServiceConfig
+from paasta_tools.chronos_tools import ChronosJobConfig
 
 
 def test_get_branches_for_service():
     fake_dir = '/mail/var/tea'
     fake_srv = 'boba'
-    fake_fnames = ['marathon-quail', 'marathon-trac', 'chronos-other', 'fake-file']
-    fake_branches = [set(['badboy']), set(['red', 'green']), set(['blue', 'orange']),  set(['white', 'black'])]
-    expected = set(['red', 'green', 'blue', 'orange', 'white', 'black'])
+    fake_branches = ['red', 'green', 'blue', 'orange', 'white', 'black']
+    expected = set(['red', 'green', 'blue', 'orange', 'white', 'black', 'cluster_a.chronos_job',
+                    'cluster_b.chronos_job', 'cluster_c.chronos_job'])
     with contextlib.nested(
-        mock.patch('os.path.join', return_value='no_sir'),
-        mock.patch('os.listdir', return_value=fake_fnames),
-        mock.patch('generate_deployments_for_service.get_branches_from_config_file',
-                   side_effect=lambda a, b: fake_branches.pop())
+        mock.patch('generate_deployments_for_service.list_clusters',
+                   return_value=['cluster_a', 'cluster_b', 'cluster_c']),
+        mock.patch('generate_deployments_for_service.get_service_instance_list',
+                   side_effect=lambda service, cluster, instance_type:
+                       [('', 'main_example'), ('', 'canary_example')] if instance_type == 'marathon'
+                       else [('', 'chronos_job')]),
+        mock.patch('generate_deployments_for_service.load_marathon_service_config',
+                   side_effect=lambda service, instance, cluster, soa_dir: MarathonServiceConfig(
+                       service=service,
+                       cluster=cluster,
+                       instance=instance,
+                       config_dict={'deploy_group': fake_branches.pop()},
+                       branch_dict={},
+                   )),
+        mock.patch('generate_deployments_for_service.load_chronos_job_config',
+                   side_effect=lambda service, instance, cluster, soa_dir: ChronosJobConfig(
+                       service=service,
+                       cluster=cluster,
+                       instance=instance,
+                       config_dict={},
+                       branch_dict={},
+                   )),
     ) as (
-        join_patch,
-        listdir_patch,
-        get_branches_patch
+        list_clusters_patch,
+        get_service_instance_list_patch,
+        load_marathon_config_patch,
+        load_chronos_config_patch,
     ):
         actual = generate_deployments_for_service.get_branches_for_service(fake_dir, fake_srv)
         assert expected == actual
-        join_patch.assert_called_once_with(fake_dir, fake_srv)
-        listdir_patch.assert_called_once_with('no_sir')
-        get_branches_patch.assert_any_call('no_sir', 'marathon-quail')
-        get_branches_patch.assert_any_call('no_sir', 'marathon-trac')
-        get_branches_patch.assert_any_call('no_sir', 'chronos-other')
-        assert get_branches_patch.call_count == 3
+        # three clusters each having a canary and main marathon instance equals six instances total
+        assert load_marathon_config_patch.call_count == 6
+        # three clusters each having one chronos job equals six instances total
+        assert load_chronos_config_patch.call_count == 3
 
 
 def test_get_branch_mappings():