You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@aurora.apache.org by ma...@apache.org on 2014/01/09 01:46:44 UTC
git commit: Client quota checks. Part 1: client side changes.
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 <mk...@twitter.com>
Authored: Wed Jan 8 16:44:59 2014 -0800
Committer: Maxim Khutornenko <mk...@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())