def test_pass_with_no_consumed(self):
allocated = ResourceAggregate(numCpus=50.0, ramMb=1000, diskMb=3000)
released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=500))
acquired = CapacityRequest(ResourceAggregate(numCpus=15.0, ramMb=300, diskMb=800))
self.mock_get_quota(allocated, None)
self.assert_result(True, released, acquired)
def test_increase_quota(self):
"""Tests successful execution of the increase_quota command."""
mock_options = self.setup_mock_options()
with contextlib.nested(
patch('twitter.common.app.get_options',
return_value=mock_options),
patch('apache.aurora.client.commands.admin.AuroraClientAPI',
new=Mock(spec=AuroraClientAPI)),
patch('apache.aurora.client.commands.admin.CLUSTERS',
new=self.TEST_CLUSTERS)) as (_, api, _):
role = 'test_role'
api.return_value.get_quota.return_value = self.create_response(
ResourceAggregate(20.0, 4000, 6000),
ResourceAggregate(15.0, 2000, 3000),
ResourceAggregate(6.0, 200, 600),
)
api.return_value.set_quota.return_value = self.create_simple_success_response(
)
increase_quota([self.TEST_CLUSTER, role, '4.0', '1MB', '1MB'])
api.return_value.set_quota.assert_called_with(
role, 24.0, 4001, 6001)
assert type(api.return_value.set_quota.call_args[0][1]) == type(
float())
assert type(api.return_value.set_quota.call_args[0][2]) == type(
int())
assert type(api.return_value.set_quota.call_args[0][3]) == type(
int())
def setup_mock_quota_call_with_consumption(cls, mock_context):
api = mock_context.get_api('west')
response = cls.create_simple_success_response()
response.result = Result(getQuotaResult=GetQuotaResult(
quota=ResourceAggregate(resources=frozenset([
Resource(numCpus=5),
Resource(ramMb=20480),
Resource(diskMb=40960)
])),
prodSharedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=1),
Resource(ramMb=512),
Resource(diskMb=1024)
])),
prodDedicatedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=2),
Resource(ramMb=1024),
Resource(diskMb=2048)
])),
nonProdSharedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=3),
Resource(ramMb=2048),
Resource(diskMb=4096)
])),
nonProdDedicatedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=4),
Resource(ramMb=4096),
Resource(diskMb=8192)
])),
))
api.get_quota.return_value = response
def test_increase_quota(self):
"""Tests successful execution of the increase_quota command."""
mock_options = self.setup_mock_options()
with contextlib.nested(
patch('twitter.common.app.get_options',
return_value=mock_options),
patch('apache.aurora.admin.admin.make_admin_client',
return_value=create_autospec(spec=AuroraClientAPI)),
patch('apache.aurora.admin.admin.CLUSTERS',
new=self.TEST_CLUSTERS)) as (_, mock_make_admin_client,
_):
api = mock_make_admin_client.return_value
role = 'test_role'
api.get_quota.return_value = self.create_response(
ResourceAggregate(20.0, 4000, 6000),
ResourceAggregate(15.0, 2000, 3000),
ResourceAggregate(6.0, 200, 600),
)
api.set_quota.return_value = self.create_simple_success_response()
increase_quota([self.TEST_CLUSTER, role, '4.0', '1MB', '1MB'])
api.set_quota.assert_called_with(role, 24.0, 4001, 6001)
assert isinstance(api.set_quota.call_args[0][1], float)
assert isinstance(api.set_quota.call_args[0][2], int)
assert isinstance(api.set_quota.call_args[0][3], int)
def test_pass_due_to_released(self):
allocated = ResourceAggregate(numCpus=50.0, ramMb=1000, diskMb=3000)
consumed = ResourceAggregate(numCpus=45.0, ramMb=900, diskMb=2900)
released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=100))
acquired = CapacityRequest(ResourceAggregate(numCpus=10.0, ramMb=200, diskMb=200))
self.mock_get_quota(allocated, consumed)
self.assert_result(True, released, acquired)
def test_fail(self):
allocated = ResourceAggregate(numCpus=50.0, ramMb=1000, diskMb=3000)
consumed = ResourceAggregate(numCpus=25.0, ramMb=500, diskMb=2000)
released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=500))
acquired = CapacityRequest(ResourceAggregate(numCpus=35.0, ramMb=300, diskMb=800))
self.mock_get_quota(allocated, consumed)
self.assert_result(True, released, acquired, ResponseCode.INVALID_REQUEST)
def setup_mock_quota_call_with_consumption(cls, mock_context):
api = mock_context.get_api('west')
response = cls.create_simple_success_response()
response.result = Result(getQuotaResult=GetQuotaResult(
quota=ResourceAggregate(numCpus=5, ramMb=20480, diskMb=40960),
prodConsumption=ResourceAggregate(
numCpus=1, ramMb=1024, diskMb=2048),
nonProdConsumption=ResourceAggregate(
numCpus=1, ramMb=1024, diskMb=2048),
))
api.get_quota.return_value = response
def test_additional_quota_out(self, mock_print_quota):
allocated = ResourceAggregate(numCpus=50.0, ramMb=1000, diskMb=3000)
consumed = ResourceAggregate(numCpus=45.0, ramMb=900, diskMb=2900)
released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=100))
acquired = CapacityRequest(ResourceAggregate(numCpus=11.0, ramMb=220, diskMb=200))
additional = ResourceAggregate(numCpus=1.0, ramMb=20, diskMb=0)
self.mock_get_quota(allocated, consumed)
self.assert_result(True, released, acquired, ResponseCode.INVALID_REQUEST)
assert mock_print_quota.mock_calls[:4] == [
call(allocated, 'Total allocated quota', self._role),
call(consumed, 'Consumed quota', self._role),
call((acquired - released).quota(), 'Requested', self._name),
call(additional, 'Additional quota required', self._role)
]
def test_setQuota(self):
self.mock_thrift_client.setQuota(
IgnoreArg(),
IsA(ResourceAggregate),
SESSION).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().setQuota(ROLE, ResourceAggregate())
def expect_quota_check(self,
num_released,
num_acquired,
response_code=ResponseCode.OK,
prod=True):
released = CapacityRequest(ResourceAggregate(
numCpus=num_released * self._num_cpus,
ramMb=num_released * self._num_ram,
diskMb=num_released * self._num_disk))
acquired = CapacityRequest(ResourceAggregate(
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(make_response(response_code))
def test_backfill_quota(self):
quota = ResourceAggregate(numCpus=1.0, ramMb=2, diskMb=3)
assert ResourceManager.resource_details_from_quota(quota) == [
ResourceDetails(ResourceType.CPUS, 1.0),
ResourceDetails(ResourceType.RAM_MB, 2),
ResourceDetails(ResourceType.DISK_MB, 3)
]
def _op(self, op, other):
if not isinstance(other, CapacityRequest):
return self
return CapacityRequest(
ResourceAggregate(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 set_quota(self, role, cpu, ram, disk):
log.info("Setting quota for user:%s cpu:%f ram:%d disk: %d"
% (role, cpu, ram, disk))
return self._scheduler_proxy.setQuota(
role,
ResourceAggregate(cpu, ram, disk, frozenset([
Resource(numCpus=cpu),
Resource(ramMb=ram),
Resource(diskMb=disk)])))
def invert_or_reset(self):
"""Inverts negative resource and resets positive resource as zero."""
def invert_or_reset(val):
return math.fabs(val) if val < 0 else 0
return CapacityRequest(
ResourceAggregate(numCpus=invert_or_reset(self._quota.numCpus),
ramMb=invert_or_reset(self._quota.ramMb),
diskMb=invert_or_reset(self._quota.diskMb)))
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,
messageDEPRECATED='test')
released = CapacityRequest(
ResourceAggregate(numCpus=num_released * self._num_cpus,
ramMb=num_released * self._num_ram,
diskMb=num_released * self._num_disk))
acquired = CapacityRequest(
ResourceAggregate(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 __init__(self, quota=None): self._quota = quota or ResourceAggregate(numCpus=0.0, ramMb=0, diskMb=0)
def from_task(cls, task): return cls(ResourceAggregate(numCpus=task.numCpus, ramMb=task.ramMb, diskMb=task.diskMb))
def test_setQuota(self):
self.mock_thrift_client.setQuota(IgnoreArg(), IsA(ResourceAggregate),
IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().setQuota(ROLE, ResourceAggregate())
def test_right_add(self): released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=100)) # Testing __radd__ which is called when object on the ride side isn't our CapacityRequest # instance. out = 1 + released self.assertEqual(out, released)
def test_op_not_cap_request(self):
released = CapacityRequest(ResourceAggregate(numCpus=5.0, ramMb=100, diskMb=100))
# Should return self so a noop.
out = released._op('not_a_real_op', 'not_a_CapacityRequest_obj')
self.assertEqual(out, released)
def set_quota(self, role, cpu, ram, disk):
log.info("Setting quota for user:%s cpu:%f ram:%d disk: %d"
% (role, cpu, ram, disk))
return self._scheduler_proxy.setQuota(role, ResourceAggregate(cpu, ram, disk))
