aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject git commit: Client quota checks. Part 1: client side changes.
Date Thu, 09 Jan 2014 00:46:44 GMT
Updated Branches:
  refs/heads/master d801596fa -> 7e87588c1


Client quota checks. Part 1: client side changes.

Reviewed at https://reviews.apache.org/r/16444/


Project: http://git-wip-us.apache.org/repos/asf/incubator-aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-aurora/commit/7e87588c
Tree: http://git-wip-us.apache.org/repos/asf/incubator-aurora/tree/7e87588c
Diff: http://git-wip-us.apache.org/repos/asf/incubator-aurora/diff/7e87588c

Branch: refs/heads/master
Commit: 7e87588c154c0faf0cb99baa46749ceb611abf43
Parents: d801596
Author: Maxim Khutornenko <mkhutornenko@twitter.com>
Authored: Wed Jan 8 16:44:59 2014 -0800
Committer: Maxim Khutornenko <mkhutornenko@twitter.com>
Committed: Wed Jan 8 16:44:59 2014 -0800

----------------------------------------------------------------------
 src/main/python/apache/aurora/client/api/BUILD  | 11 +++
 .../apache/aurora/client/api/quota_check.py     | 97 ++++++++++++++++++++
 .../python/apache/aurora/client/api/updater.py  | 39 +++++++-
 .../python/apache/aurora/client/commands/BUILD  |  1 +
 .../apache/aurora/client/commands/core.py       | 15 ++-
 .../thrift/org/apache/aurora/gen/api.thrift     |  3 +-
 src/test/python/apache/aurora/client/api/BUILD  | 10 ++
 .../aurora/client/api/test_quota_check.py       | 95 +++++++++++++++++++
 .../apache/aurora/client/api/test_updater.py    | 70 +++++++++++++-
 .../aurora/client/commands/test_update.py       | 10 +-
 10 files changed, 333 insertions(+), 18 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/python/apache/aurora/client/api/BUILD
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/BUILD b/src/main/python/apache/aurora/client/api/BUILD
index 9af74e9..f0f57e0 100644
--- a/src/main/python/apache/aurora/client/api/BUILD
+++ b/src/main/python/apache/aurora/client/api/BUILD
@@ -84,12 +84,23 @@ python_library(
 )
 
 python_library(
+  name = 'quota_check',
+  sources = ['quota_check.py'],
+  dependencies = [
+    pants(':scheduler_client'),
+    pants('aurora/twitterdeps/src/python/twitter/common/log'),
+    pants('src/main/thrift/org/apache/aurora/gen:py-thrift'),
+  ]
+)
+
+python_library(
   name = 'updater',
   sources = ['updater.py'],
   dependencies = [
     pants(':scheduler_client'),
     pants(':instance_watcher'),
     pants(':updater_util'),
+    pants(':quota_check'),
     pants('aurora/twitterdeps/src/python/twitter/common/log'),
     pants('src/main/thrift/org/apache/aurora/gen:py-thrift'),
   ]

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/python/apache/aurora/client/api/quota_check.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/quota_check.py b/src/main/python/apache/aurora/client/api/quota_check.py
new file mode 100644
index 0000000..00f76d9
--- /dev/null
+++ b/src/main/python/apache/aurora/client/api/quota_check.py
@@ -0,0 +1,97 @@
+import operator
+
+from copy import deepcopy
+
+from twitter.common import log
+
+from gen.apache.aurora.ttypes import Quota, Response, ResponseCode
+
+
+class CapacityRequest(object):
+  """Facilitates Quota manipulations."""
+
+  @classmethod
+  def from_task(cls, task):
+    return cls(Quota(numCpus=task.numCpus, ramMb=task.ramMb, diskMb=task.diskMb))
+
+  def __init__(self, quota=None):
+    self._quota = quota or Quota(numCpus=0.0, ramMb=0, diskMb=0)
+
+  def __add__(self, other):
+    return self._op(operator.__add__, other)
+
+  def __radd__(self, other):
+    return self._op(operator.__add__, other)
+
+  def __sub__(self, other):
+    return self._op(operator.__sub__, other)
+
+  def __eq__(self, other):
+    return self._quota == other._quota
+
+  def _op(self, op, other):
+    if not isinstance(other, CapacityRequest):
+      return self
+
+    return CapacityRequest(
+        Quota(numCpus=op(self._quota.numCpus, other._quota.numCpus),
+              ramMb=op(self._quota.ramMb, other._quota.ramMb),
+              diskMb=op(self._quota.diskMb, other._quota.diskMb)))
+
+  def valid(self):
+    return self._quota.numCpus >= 0.0 and self._quota.ramMb >= 0 and self._quota.diskMb
>= 0
+
+  def quota(self):
+    return deepcopy(self._quota)
+
+
+class QuotaCheck(object):
+  """Performs quota checks for the provided job/task configurations."""
+
+  def __init__(self, scheduler):
+    self._scheduler = scheduler
+
+  def validate_quota_from_requested(self, job_key, production, released, acquired):
+    """Validates requested change will not exceed the available quota.
+
+    Arguments:
+    job_key -- job key.
+    production -- production flag.
+    released -- CapacityRequest to be released (in case of job update).
+    acquired -- CapacityRequest to be acquired.
+
+    Returns: ResponseCode.OK if check is successful.
+    """
+    resp_ok = Response(responseCode=ResponseCode.OK, message='Quota check successful.')
+    if not production:
+      return resp_ok
+
+    resp = self._scheduler.getQuota(job_key.role)
+    if resp.responseCode != ResponseCode.OK:
+      log.error('Failed to get quota from scheduler: %s' % resp.message)
+      return resp
+
+    allocated = CapacityRequest(resp.result.getQuotaResult.quota)
+    consumed = CapacityRequest(resp.result.getQuotaResult.consumed)
+    requested = acquired - released
+    effective = allocated - consumed - requested
+
+    if not effective.valid():
+      log.info('Not enough quota to create/update job.')
+      print_quota(allocated.quota(), 'Total allocated quota', job_key.role)
+      print_quota(consumed.quota(), 'Consumed quota', job_key.role)
+      print_quota(requested.quota(), 'Requested', job_key.name)
+      return Response(responseCode=ResponseCode.INVALID_REQUEST, message='Failed quota check.')
+
+    return resp_ok
+
+
+def print_quota(quota, msg, subj):
+  quota_fields = [
+      ('CPU', quota.numCpus),
+      ('RAM', '%f GB' % (float(quota.ramMb) / 1024)),
+      ('Disk', '%f GB' % (float(quota.diskMb) / 1024))
+  ]
+  log.info('%s for %s:\n\t%s' %
+           (msg, subj, '\n\t'.join(['%s\t%s' % (k, v) for (k, v) in quota_fields])))
+

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/python/apache/aurora/client/api/updater.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/updater.py b/src/main/python/apache/aurora/client/api/updater.py
index c885340..c915a7d 100644
--- a/src/main/python/apache/aurora/client/api/updater.py
+++ b/src/main/python/apache/aurora/client/api/updater.py
@@ -14,14 +14,16 @@ from gen.apache.aurora.ttypes import (
     Lock,
     LockKey,
     LockValidation,
+    Quota,
     Response,
     ResponseCode,
     TaskQuery,
 )
 
-from .updater_util import FailureThreshold, UpdaterConfig
 from .instance_watcher import InstanceWatcher
+from .quota_check import CapacityRequest, QuotaCheck
 from .scheduler_client import SchedulerProxy
+from .updater_util import FailureThreshold, UpdaterConfig
 
 from thrift.protocol import TJSONProtocol
 from thrift.TSerialization import serialize
@@ -38,11 +40,17 @@ class Updater(object):
       ['remote_config_map', 'local_config_map', 'instances_to_process']
   )
 
-  def __init__(self, config, health_check_interval_seconds, scheduler=None, instance_watcher=None):
+  def __init__(self,
+               config,
+               health_check_interval_seconds,
+               scheduler=None,
+               instance_watcher=None,
+               quota_check=None):
     self._config = config
     self._job_key = JobKey(role=config.role(), environment=config.environment(), name=config.name())
     self._health_check_interval_seconds = health_check_interval_seconds
     self._scheduler = scheduler or SchedulerProxy(config.cluster())
+    self._quota_check = quota_check or QuotaCheck(self._scheduler)
     try:
       self._update_config = UpdaterConfig(**config.update_config().get())
     except ValueError as e:
@@ -283,6 +291,32 @@ class Updater(object):
     self._check_and_log_response(resp)
     return instance_ids
 
+  def _validate_quota(self, instance_configs):
+    """Validates job update will not exceed quota for production tasks.
+    Arguments:
+    instance_configs -- InstanceConfig with update details.
+
+    Returns Response.OK if quota check was successful.
+    """
+    instance_operation = self.OperationConfigs(
+      from_config=instance_configs.remote_config_map,
+      to_config=instance_configs.local_config_map
+    )
+
+    def _aggregate_quota(ops_list, config_map):
+      return sum(CapacityRequest.from_task(config_map[instance])
+                    for instance in ops_list) or CapacityRequest()
+
+    to_kill, to_add = self._create_kill_add_lists(
+        instance_configs.instances_to_process,
+        instance_operation)
+
+    return self._quota_check.validate_quota_from_requested(
+        self._job_key,
+        self._config.job().taskConfig.production,
+        _aggregate_quota(to_kill, instance_operation.from_config),
+        _aggregate_quota(to_add, instance_operation.to_config))
+
   def _get_update_instructions(self, instances=None):
     """Loads, validates and populates update working set.
 
@@ -392,6 +426,7 @@ class Updater(object):
       else:
         try:
           instance_configs = self._get_update_instructions(instances)
+          self._check_and_log_response(self._validate_quota(instance_configs))
         except self.Error as e:
           # Safe to release the lock acquired above as no job mutation has happened yet.
           self._finish()

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/python/apache/aurora/client/commands/BUILD
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/BUILD b/src/main/python/apache/aurora/client/commands/BUILD
index 556d00a..6b4660b 100644
--- a/src/main/python/apache/aurora/client/commands/BUILD
+++ b/src/main/python/apache/aurora/client/commands/BUILD
@@ -33,6 +33,7 @@ python_library(
     pants('src/main/python/apache/aurora/client/api:command_runner'),
     pants('src/main/python/apache/aurora/client/api:disambiguator'),
     pants('src/main/python/apache/aurora/client/api:job_monitor'),
+    pants('src/main/python/apache/aurora/client/api:quota_check'),
     pants('src/main/python/apache/aurora/client/api:updater'),
     pants('src/main/python/apache/aurora/client/hooks'),
     pants('src/main/python/apache/aurora/client:base'),

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/python/apache/aurora/client/commands/core.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/core.py b/src/main/python/apache/aurora/client/commands/core.py
index 45204b5..edc57b2 100644
--- a/src/main/python/apache/aurora/client/commands/core.py
+++ b/src/main/python/apache/aurora/client/commands/core.py
@@ -26,6 +26,7 @@ from apache.aurora.client.base import (
     synthesize_url)
 from apache.aurora.client.api.disambiguator import LiveJobDisambiguator
 from apache.aurora.client.api.job_monitor import JobMonitor
+from apache.aurora.client.api.quota_check import print_quota
 from apache.aurora.client.api.updater_util import UpdaterConfig
 from apache.aurora.client.config import get_config
 from apache.aurora.client.factory import make_client, make_client_factory
@@ -589,12 +590,8 @@ def get_quota(role):
   """
   options = app.get_options()
   resp = make_client(options.cluster).get_quota(role)
-  quota = resp.result.getQuotaResult.quota
-
-  quota_fields = [
-    ('CPU', quota.numCpus),
-    ('RAM', '%f GB' % (float(quota.ramMb) / 1024)),
-    ('Disk', '%f GB' % (float(quota.diskMb) / 1024))
-  ]
-  log.info('Quota for %s:\n\t%s' %
-           (role, '\n\t'.join(['%s\t%s' % (k, v) for (k, v) in quota_fields])))
+
+  print_quota(resp.result.getQuotaResult.quota, 'Total allocated quota', role)
+
+  if resp.result.getQuotaResult.consumed:
+    print_quota(resp.result.getQuotaResult.consumed, 'Consumed quota', role)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/main/thrift/org/apache/aurora/gen/api.thrift
----------------------------------------------------------------------
diff --git a/src/main/thrift/org/apache/aurora/gen/api.thrift b/src/main/thrift/org/apache/aurora/gen/api.thrift
index 480b8f4..33c70df 100644
--- a/src/main/thrift/org/apache/aurora/gen/api.thrift
+++ b/src/main/thrift/org/apache/aurora/gen/api.thrift
@@ -212,7 +212,8 @@ struct PopulateJobResult {
 }
 
 struct GetQuotaResult {
-  1: Quota quota
+  1: Quota quota              // Total allocated quota.
+  2: optional Quota consumed  // Amount of quota already consumed by a role.
 }
 
 // Wraps return results for the acquireLock API.

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/test/python/apache/aurora/client/api/BUILD
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/BUILD b/src/test/python/apache/aurora/client/api/BUILD
index c532675..4c4bcca 100644
--- a/src/test/python/apache/aurora/client/api/BUILD
+++ b/src/test/python/apache/aurora/client/api/BUILD
@@ -6,6 +6,7 @@ python_test_suite(name = 'all',
     pants(':scheduler_client'),
     pants(':instance_watcher'),
     pants(':updater'),
+    pants(':quota_check'),
   ],
 )
 
@@ -59,6 +60,15 @@ python_tests(name = 'instance_watcher',
   ]
 )
 
+python_tests(name = 'quota_check',
+  sources = ['test_quota_check.py'],
+  dependencies = [
+    pants('src/main/python/apache/aurora/BUILD.thirdparty:mock'),
+    pants('src/main/python/apache/aurora/client/api:quota_check'),
+    pants('src/main/thrift/org/apache/aurora/gen:py-thrift'),
+  ]
+)
+
 python_tests(name = 'updater',
   sources = ['test_updater.py'],
   dependencies = [

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/test/python/apache/aurora/client/api/test_quota_check.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_quota_check.py b/src/test/python/apache/aurora/client/api/test_quota_check.py
new file mode 100644
index 0000000..c1fa6d0
--- /dev/null
+++ b/src/test/python/apache/aurora/client/api/test_quota_check.py
@@ -0,0 +1,95 @@
+import unittest
+
+from copy import deepcopy
+
+from apache.aurora.client.api.quota_check import CapacityRequest, QuotaCheck
+
+from gen.apache.aurora.AuroraSchedulerManager import Client as scheduler_client
+from gen.apache.aurora.ttypes import (
+    GetQuotaResult,
+    JobKey,
+    Quota,
+    Response,
+    ResponseCode,
+    Result)
+
+from mock import Mock
+
+
+class QuotaCheckTest(unittest.TestCase):
+  def setUp(self):
+    self._scheduler = Mock()
+    self._quota_checker = QuotaCheck(self._scheduler)
+    self._role = 'mesos'
+    self._name = 'quotajob'
+    self._env = 'test'
+    self._job_key = JobKey(name=self._name, environment=self._env, role=self._role)
+
+  def mock_get_quota(self, allocated, consumed, response_code=None):
+    response_code = ResponseCode.OK if response_code is None else response_code
+
+    resp = Response(responseCode=response_code, message='test')
+    resp.result = Result(
+        getQuotaResult=GetQuotaResult(quota=deepcopy(allocated), consumed=deepcopy(consumed)))
+    self._scheduler.getQuota.return_value = resp
+
+  def assert_result(self, prod, released, acquired, expected_code=None):
+    expected_code = ResponseCode.OK if expected_code is None else expected_code
+    resp = self._quota_checker.validate_quota_from_requested(
+        self._job_key,
+        prod,
+        released,
+        acquired)
+    assert expected_code == resp.responseCode, (
+      'Expected response:%s Actual response:%s' % (expected_code, resp.responseCode))
+    if prod:
+      self._scheduler.getQuota.assert_called_once_with(self._role)
+    else:
+      assert not self._scheduler.getQuota.called, 'Scheduler.getQuota() unexpected call.'
+
+  def test_pass(self):
+    allocated = Quota(numCpus=50.0, ramMb=1000, diskMb=3000)
+    consumed = Quota(numCpus=25.0, ramMb=500, diskMb=2000)
+    released = CapacityRequest(Quota(numCpus=5.0, ramMb=100, diskMb=500))
+    acquired = CapacityRequest(Quota(numCpus=15.0, ramMb=300, diskMb=800))
+
+    self.mock_get_quota(allocated, consumed)
+    self.assert_result(True, released, acquired)
+
+  def test_pass_with_no_consumed(self):
+    allocated = Quota(numCpus=50.0, ramMb=1000, diskMb=3000)
+    released = CapacityRequest(Quota(numCpus=5.0, ramMb=100, diskMb=500))
+    acquired = CapacityRequest(Quota(numCpus=15.0, ramMb=300, diskMb=800))
+
+    self.mock_get_quota(allocated, None)
+    self.assert_result(True, released, acquired)
+
+  def test_pass_due_to_released(self):
+    allocated = Quota(numCpus=50.0, ramMb=1000, diskMb=3000)
+    consumed = Quota(numCpus=45.0, ramMb=900, diskMb=2900)
+    released = CapacityRequest(Quota(numCpus=5.0, ramMb=100, diskMb=100))
+    acquired = CapacityRequest(Quota(numCpus=10.0, ramMb=200, diskMb=200))
+
+    self.mock_get_quota(allocated, consumed)
+    self.assert_result(True, released, acquired)
+
+  def test_skipped(self):
+    self.assert_result(False, None, None)
+
+  def test_fail(self):
+    allocated = Quota(numCpus=50.0, ramMb=1000, diskMb=3000)
+    consumed = Quota(numCpus=25.0, ramMb=500, diskMb=2000)
+    released = CapacityRequest(Quota(numCpus=5.0, ramMb=100, diskMb=500))
+    acquired = CapacityRequest(Quota(numCpus=35.0, ramMb=300, diskMb=800))
+
+    self.mock_get_quota(allocated, consumed)
+    self.assert_result(True, released, acquired, ResponseCode.INVALID_REQUEST)
+
+  def test_fail_scheduler_call(self):
+    allocated = Quota(numCpus=50.0, ramMb=1000, diskMb=3000)
+    consumed = Quota(numCpus=25.0, ramMb=500, diskMb=2000)
+    released = CapacityRequest(Quota(numCpus=5.0, ramMb=100, diskMb=500))
+    acquired = CapacityRequest(Quota(numCpus=1.0, ramMb=100, diskMb=100))
+
+    self.mock_get_quota(allocated, consumed, response_code=ResponseCode.INVALID_REQUEST)
+    self.assert_result(True, released, acquired, ResponseCode.INVALID_REQUEST)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/test/python/apache/aurora/client/api/test_updater.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_updater.py b/src/test/python/apache/aurora/client/api/test_updater.py
index 96371aa..2e8c8a8 100644
--- a/src/test/python/apache/aurora/client/api/test_updater.py
+++ b/src/test/python/apache/aurora/client/api/test_updater.py
@@ -3,6 +3,7 @@ from os import environ
 from unittest import TestCase
 
 from apache.aurora.client.api.instance_watcher import InstanceWatcher
+from apache.aurora.client.api.quota_check import CapacityRequest, QuotaCheck
 from apache.aurora.client.api.updater import Updater
 from apache.aurora.client.fake_scheduler_proxy import FakeSchedulerProxy
 
@@ -24,6 +25,7 @@ from gen.apache.aurora.ttypes import (
   LockValidation,
   Package,
   PopulateJobResult,
+  Quota,
   Response,
   ResponseCode,
   Result,
@@ -101,19 +103,30 @@ class UpdaterTest(TestCase):
     self._instance_watcher = MockObject(InstanceWatcher)
     self._scheduler = MockObject(scheduler_client)
     self._scheduler_proxy = FakeSchedulerProxy('test-cluster', self._scheduler, self._session_key)
+    self._quota_check = MockObject(QuotaCheck)
     self.init_updater(deepcopy(self.UPDATE_CONFIG))
+    self._num_cpus = 1.0
+    self._num_ram = 1
+    self._num_disk = 1
 
   def replay_mocks(self):
     Replay(self._scheduler)
     Replay(self._instance_watcher)
+    Replay(self._quota_check)
 
   def verify_mocks(self):
     Verify(self._scheduler)
     Verify(self._instance_watcher)
+    Verify(self._quota_check)
 
   def init_updater(self, update_config):
     self._config = FakeConfig(self._role, self._name, self._env, update_config)
-    self._updater = Updater(self._config, 3, self._scheduler_proxy, self._instance_watcher)
+    self._updater = Updater(
+        self._config,
+        3,
+        self._scheduler_proxy,
+        self._instance_watcher,
+        self._quota_check)
 
   def expect_watch_instances(self, instance_ids, failed_instances=[]):
     self._instance_watcher.watch(instance_ids).AndReturn(set(failed_instances))
@@ -188,14 +201,29 @@ class UpdaterTest(TestCase):
         LockValidation.CHECKED,
         self._session_key).AndReturn(response)
 
+  def expect_quota_check(self, num_released, num_acquired, response_code=None, prod=True):
+    response_code = ResponseCode.OK if response_code is None else response_code
+    response = Response(responseCode=response_code, message='test')
+    released = CapacityRequest(Quota(
+        numCpus=num_released * self._num_cpus,
+        ramMb=num_released * self._num_ram,
+        diskMb=num_released * self._num_disk))
+    acquired = CapacityRequest(Quota(
+      numCpus=num_acquired * self._num_cpus,
+      ramMb=num_acquired * self._num_ram,
+      diskMb=num_acquired * self._num_disk))
+
+    self._quota_check.validate_quota_from_requested(
+        self._job_key, prod, released, acquired).AndReturn(response)
+
   def make_task_configs(self, count=1):
     return [TaskConfig(
         owner=Identity(role=self._job_key.role),
         environment=self._job_key.environment,
         jobName=self._job_key.name,
-        numCpus=6.0,
-        ramMb=1024,
-        diskMb=2048,
+        numCpus=self._num_cpus,
+        ramMb=self._num_ram,
+        diskMb=self._num_disk,
         priority=0,
         maxTaskFailures=1,
         production=True,
@@ -210,7 +238,7 @@ class UpdaterTest(TestCase):
         key=self._job_key,
         owner=Identity(role=self._job_key.role),
         cronSchedule=cron_schedule,
-        taskConfig=task_config,
+        taskConfig=deepcopy(task_config),
         instanceCount=instance_count
     )
 
@@ -231,6 +259,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(0, 4)
     self.expect_add([3, 4, 5], new_config)
     self.expect_watch_instances([3, 4, 5])
     self.expect_add([6], new_config)
@@ -241,6 +270,22 @@ class UpdaterTest(TestCase):
     self.update_and_expect_ok()
     self.verify_mocks()
 
+  def test_grow_fails_quota_check(self):
+    """Adds instances to the existing job fails due to not enough quota."""
+    old_configs = self.make_task_configs(3)
+    new_config = old_configs[0]
+    job_config = self.make_job_config(new_config, 7)
+    self._config.job_config = job_config
+    self.expect_start()
+    self.expect_get_tasks(old_configs)
+    self.expect_populate(job_config)
+    self.expect_quota_check(0, 4, response_code=ResponseCode.INVALID_REQUEST)
+    self.expect_finish()
+    self.replay_mocks()
+
+    self.update_and_expect_response(expected_code=ResponseCode.ERROR)
+    self.verify_mocks()
+
   def test_shrink(self):
     """Reduces the number of instances of the job."""
     old_configs = self.make_task_configs(10)
@@ -250,6 +295,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(7, 0)
     self.expect_kill([3, 4, 5])
     self.expect_kill([6, 7, 8])
     self.expect_kill([9])
@@ -269,6 +315,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(3, 7)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2])
@@ -292,6 +339,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(10, 1)
     self.expect_kill([0, 1, 2])
     self.expect_add([0], new_config)
     self.expect_watch_instances([0])
@@ -314,6 +362,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(5, 5)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2])
@@ -335,6 +384,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(0, 2)
     self.expect_add([3, 4], new_config)
     self.expect_watch_instances([3, 4])
     self.expect_finish()
@@ -352,6 +402,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(6, 0)
     self.expect_kill([4, 5, 6])
     self.expect_kill([7, 8, 9])
     self.expect_finish()
@@ -370,6 +421,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(3, 3)
     self.expect_kill([2, 3, 4])
     self.expect_add([2, 3, 4], new_config)
     self.expect_watch_instances([2, 3, 4])
@@ -389,6 +441,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs, [2, 3])
     self.expect_populate(job_config)
+    self.expect_quota_check(0, 2)
     self.expect_add([2, 3], new_config)
     self.expect_watch_instances([2, 3])
     self.expect_finish()
@@ -406,6 +459,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(0, 0)
     self.expect_finish()
     self.replay_mocks()
 
@@ -426,6 +480,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(10, 10)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2], failed_instances=[0, 1, 2])
@@ -454,6 +509,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(5, 5)
     self.expect_kill([0])
     self.expect_add([0], new_config)
     self.expect_watch_instances([0])
@@ -494,6 +550,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(6, 6)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2], failed_instances=[0, 1, 2])
@@ -573,6 +630,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(5, 5)
     self.expect_kill([0, 1, 2], response_code=ResponseCode.INVALID_REQUEST)
     self.replay_mocks()
 
@@ -623,6 +681,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(10, 10)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2], failed_instances=[0])
@@ -657,6 +716,7 @@ class UpdaterTest(TestCase):
     self.expect_start()
     self.expect_get_tasks(old_configs)
     self.expect_populate(job_config)
+    self.expect_quota_check(5, 5)
     self.expect_kill([0, 1, 2])
     self.expect_add([0, 1, 2], new_config)
     self.expect_watch_instances([0, 1, 2], failed_instances=[0])

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/7e87588c/src/test/python/apache/aurora/client/commands/test_update.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_update.py b/src/test/python/apache/aurora/client/commands/test_update.py
index 5d1c197..69da111 100644
--- a/src/test/python/apache/aurora/client/commands/test_update.py
+++ b/src/test/python/apache/aurora/client/commands/test_update.py
@@ -8,6 +8,7 @@ from apache.aurora.client.commands.core import update
 from apache.aurora.client.commands.util import AuroraClientCommandTest
 from apache.aurora.client.api.updater import Updater
 from apache.aurora.client.api.health_check import InstanceWatcherHealthCheck, Retriable
+from apache.aurora.client.api.quota_check import QuotaCheck
 from apache.aurora.client.hooks.hooked_api import HookedAuroraClientAPI
 from apache.aurora.config import AuroraConfig
 from twitter.common.contextutil import temporary_file
@@ -187,12 +188,18 @@ class TestUpdateCommand(AuroraClientCommandTest):
     mock_health_check.health.return_value = Retriable.alive()
     return mock_health_check
 
+  @classmethod
+  def setup_quota_check(cls):
+    mock_quota_check = Mock(spec=QuotaCheck)
+    mock_quota_check.validate_quota_from_requested.return_value = cls.create_simple_success_response()
+
   def test_updater_simple(self):
     # Test the client-side updater logic in its simplest case: everything succeeds, and no
rolling
     # updates.
     mock_options = self.setup_mock_options()
     (mock_api, mock_scheduler) = self.create_mock_api()
     mock_health_check = self.setup_health_checks(mock_api)
+    mock_quota_check = self.setup_quota_check()
 
     with contextlib.nested(
         patch('twitter.common.app.get_options', return_value=mock_options),
@@ -200,11 +207,12 @@ class TestUpdateCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
         patch('apache.aurora.client.api.instance_watcher.InstanceWatcherHealthCheck',
             return_value=mock_health_check),
+        patch('apache.aurora.client.api.quota_check.QuotaCheck', return_value=mock_quota_check),
         patch('time.time', side_effect=functools.partial(self.fake_time, self)),
         patch('time.sleep', return_value=None)
 
     ) as (options, scheduler_proxy_class, test_clusters, mock_health_check_factory,
-          time_patch, sleep_patch):
+          mock_quota_check_patch, time_patch, sleep_patch):
       self.setup_mock_scheduler_for_simple_update(mock_api)
       with temporary_file() as fp:
         fp.write(self.get_valid_config())


Mime
View raw message