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))