Allow users to set accepted_resource_roles in marathon yamls. PAASTA-…

…2455
Yelp · Jan 16, 2016 · 803b82b · 803b82b
1 parent 09bb8fa
commit 803b82b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/paasta_tools/marathon_tools.py b/paasta_tools/marathon_tools.py
@@ -312,6 +312,11 @@ def format_marathon_app_dict(self, app_id, docker_url, docker_volumes, service_n
             'cmd': self.get_cmd(),
             'args': self.get_args(),
         }
+
+        accepted_resource_roles = self.get_accepted_resource_roles()
+        if accepted_resource_roles is not None:
+            complete_config['accepted_resource_roles'] = accepted_resource_roles
+
         log.debug("Complete configuration for instance is: %s", complete_config)
         return complete_config
 
@@ -430,6 +435,9 @@ def get_bounce_health_params(self, service_namespace_config):
             default = {'check_haproxy': True}
         return self.config_dict.get('bounce_health_params', default)
 
+    def get_accepted_resource_roles(self):
+        return self.config_dict.get('accepted_resource_roles', None)
+
 
 def load_service_namespace_config(service, namespace, soa_dir=DEFAULT_SOA_DIR):
     """Attempt to read the configuration for a service's namespace in a more strict fashion.

diff --git a/tests/test_marathon_tools.py b/tests/test_marathon_tools.py
@@ -820,6 +820,7 @@ def test_format_marathon_app_dict(self):
             'health_checks': fake_healthchecks,
             'backoff_factor': 2,
             'backoff_seconds': mock.ANY,
+            'accepted_resource_roles': ['ads'],
         }
         config = marathon_tools.MarathonServiceConfig(
             'can_you_dig_it',
@@ -835,6 +836,7 @@ def test_format_marathon_app_dict(self):
                 'healthcheck_interval_seconds':  10,
                 'healthcheck_timeout_seconds':  10,
                 'healthcheck_max_consecutive_failures': 3,
+                'accepted_resource_roles': ['ads'],
             },
             {'desired_state': 'start'}
         )
@@ -1580,6 +1582,16 @@ def test_get_backoff_doesnt_devide_by_zero(self):
             "service", "instance", {'instances': 0}, {})
         assert marathon_config.get_backoff_seconds() == 1
 
+    def test_get_accepted_resource_roles_default(self):
+        marathon_config = marathon_tools.MarathonServiceConfig(
+            "service", "instance", {}, {})
+        assert marathon_config.get_accepted_resource_roles() is None
+
+    def test_get_accepted_resource_roles(self):
+        marathon_config = marathon_tools.MarathonServiceConfig(
+            "service", "instance", {"accepted_resource_roles": ["ads"]}, {})
+        assert marathon_config.get_accepted_resource_roles() == ["ads"]
+
 
 class TestServiceNamespaceConfig(object):