def test_machine_type_one_pending(self): machine_type = lease_management.MachineType( id='one-pending', leases=[ lease_management.MachineLease( client_request_id='fake-id-1', hostname='fake-host-1', request_hash='fake-hash-1', ), lease_management.MachineLease( client_request_id='fake-id-2', hostname='fake-host-2', request_hash='fake-hash-2', ), lease_management.MachineLease( client_request_id='fake-id-3', request_hash='fake-hash-3', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='one-pending', target_size=2, ) requests = lease_management._generate_lease_request_status_updates( machine_type, 'https://example.com') self.assertEqual(len(requests), 1) self.assertEqual(requests[0].request_id, 'fake-id-3')
def test_lease_denied(self): machine_type = lease_management.MachineType( id='denied', leases=[ lease_management.MachineLease(client_request_id='fake-id-1'), lease_management.MachineLease(client_request_id='fake-id-2'), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='denied', target_size=2, ) machine_type.put() responses = rpc_to_json( machine_provider.BatchedLeaseResponse(responses=[ machine_provider.LeaseResponse( client_request_id='fake-id-1', request_hash='fake-hash-1', state=machine_provider.LeaseRequestState.DENIED, ), ])) lease_management.update_leases(machine_type.key, responses) updated_machine_type = machine_type.key.get() self.assertEqual(len(updated_machine_type.leases), 1) self.assertEqual(updated_machine_type.leases[0].client_request_id, 'fake-id-2') self.failIf(updated_machine_type.leases[0].hostname) self.failIf(updated_machine_type.leases[0].lease_expiration_ts) self.failIf(updated_machine_type.leases[0].request_hash)
def test_machine_type_one_expired_lease(self): machine_type = lease_management.MachineType( id='no-leases', leases=[ lease_management.MachineLease( client_request_id='fake-id-1', request_hash='fake-hash-1', ), lease_management.MachineLease( client_request_id='fake-id-2', hostname='fake-host-2', lease_expiration_ts=utils.EPOCH, request_hash='fake-hash-2', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), pending_deletion=[ 'fake-host-3', ], request_id_base='no-leases', target_size=2, ) expired = lease_management._clean_up_expired_leases(machine_type) self.assertEqual(len(expired), 1) self.assertEqual(expired[0], 'fake-host-2') self.assertEqual(len(machine_type.leases), 1) self.assertEqual(machine_type.leases[0].client_request_id, 'fake-id-1') self.assertEqual(machine_type.leases[0].request_hash, 'fake-hash-1') self.assertEqual(len(machine_type.pending_deletion), 2) hostnames = sorted(machine_type.pending_deletion) self.assertEqual(hostnames[0], 'fake-host-2') self.assertEqual(hostnames[1], 'fake-host-3')
def test_machine_type_no_expired_leases(self): machine_type = lease_management.MachineType( id='no-leases', leases=[ lease_management.MachineLease( client_request_id='fake-id-1', request_hash='fake-hash-1', ), lease_management.MachineLease( client_request_id='fake-id-2', hostname='fake-host', lease_expiration_ts=(utils.utcnow() + datetime.timedelta(seconds=60)), request_hash='fake-hash-2', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='no-leases', target_size=2, ) expired = lease_management._clean_up_expired_leases(machine_type) self.failIf(expired) self.assertEqual(len(machine_type.leases), 2) self.failIf(machine_type.pending_deletion)
def test_dead(self): def is_dead(_self, _now): return True self.mock(bot_management.BotInfo, 'is_dead', is_dead) self.mock(lease_management, 'release', lambda *args, **kwargs: True) machine_lease = lease_management.MachineLease( bot_id='bot-id', client_request_id='req-id', hostname='bot-id', instruction_ts=utils.utcnow(), machine_type=ndb.Key(lease_management.MachineType, 'mt'), ) machine_lease.put() bot_management.bot_event( event_type='bot_leased', bot_id=machine_lease.hostname, external_ip=None, authenticated_as=None, dimensions=None, state=None, version=None, quarantined=False, maintenance_msg=None, task_id='', task_name=None, ) lease_management.check_for_connection(machine_lease) self.failIf(bot_management.get_info_key(machine_lease.bot_id).get()) self.failIf(machine_lease.key.get().client_request_id) self.failIf(machine_lease.key.get().connection_ts)
def test_lease_expired(self): machine_type = lease_management.MachineType( id='untriaged', leases=[ lease_management.MachineLease(client_request_id='fake-id-1'), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='untriaged', target_size=2, ) machine_type.put() responses = rpc_to_json( machine_provider.BatchedLeaseResponse(responses=[ machine_provider.LeaseResponse( client_request_id='fake-id-1', request_hash='fake-hash-1', state=machine_provider.LeaseRequestState.FULFILLED, ), ])) lease_management.update_leases(machine_type.key, responses) updated_machine_type = machine_type.key.get() self.failIf(updated_machine_type.leases)
def test_not_connected(self): machine_lease = lease_management.MachineLease( bot_id='bot-id', client_request_id='req-id', hostname='bot-id', instruction_ts=utils.utcnow(), machine_type=ndb.Key(lease_management.MachineType, 'mt'), ) machine_lease.put() bot_management.bot_event( event_type='bot_leased', bot_id=machine_lease.hostname, external_ip=None, authenticated_as=None, dimensions=None, state=None, version=None, quarantined=False, maintenance_msg=None, task_id='', task_name=None, ) lease_management.check_for_connection(machine_lease) self.failUnless( bot_management.get_info_key(machine_lease.bot_id).get()) self.failUnless(machine_lease.key.get().client_request_id) self.failIf(machine_lease.key.get().connection_ts)
def test_leases_fulfilled(self): machine_type = lease_management.MachineType( id='fulfilled', leases=[ lease_management.MachineLease(client_request_id='fake-id-1'), lease_management.MachineLease(client_request_id='fake-id-2'), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='fulfilled', target_size=2, ) machine_type.put() responses = rpc_to_json( machine_provider.BatchedLeaseResponse(responses=[ machine_provider.LeaseResponse( client_request_id='fake-id-1', hostname='fake-host-1', lease_expiration_ts=1, request_hash='fake-hash-1', state=machine_provider.LeaseRequestState.FULFILLED, ), machine_provider.LeaseResponse( client_request_id='fake-id-2', hostname='fake-host-2', lease_expiration_ts=2, request_hash='fake-hash-2', state=machine_provider.LeaseRequestState.FULFILLED, ), ])) lease_management.update_leases(machine_type.key, responses) updated_machine_type = machine_type.key.get() self.assertEqual(len(updated_machine_type.leases), 2) self.failUnless(updated_machine_type.leases[0].hostname) self.failUnless(updated_machine_type.leases[1].hostname) self.failUnless(updated_machine_type.leases[0].lease_expiration_ts) self.failUnless(updated_machine_type.leases[1].lease_expiration_ts) request_hashes = sorted(lease.request_hash for lease in updated_machine_type.leases) self.assertEqual(request_hashes[0], 'fake-hash-1') self.assertEqual(request_hashes[1], 'fake-hash-2')
def test_drain_one(self): key = lease_management.MachineType(target_size=0, ).put() key = lease_management.MachineLease( id='%s-0' % key.id(), machine_type=key, ).put() lease_management.drain_excess() self.assertEqual(lease_management.MachineLease.query().count(), 1) self.assertTrue(key.get().drained)
def test_machine_type_at_capacity(self): machine_type = lease_management.MachineType( id='at-capacity', leases=[ lease_management.MachineLease( client_request_id='fake-id-1', request_hash='fake-hash-1', ), lease_management.MachineLease( client_request_id='fake-id-2', request_hash='fake-hash-2', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='at-capacity', target_size=2, ) requests = lease_management._generate_lease_requests_for_new_machines( machine_type, 'https://example.com') self.failIf(requests)
def test_drain_batched(self): key = lease_management.MachineType( enabled=False, target_size=2, ).put() lease_management.MachineLease( id='%s-0' % key.id(), machine_type=key, ).put() lease_management.MachineLease( id='%s-1' % key.id(), machine_type=key, ).put() key = lease_management.MachineType( enabled=False, target_size=2, ).put() lease_management.MachineLease( id='%s-0' % key.id(), machine_type=key, ).put() lease_management.MachineLease( id='%s-1' % key.id(), machine_type=key, ).put() key = lease_management.MachineType(target_size=0, ).put() lease_management.MachineLease( id='%s-0' % key.id(), machine_type=key, ).put() # Choice of 2, 2, 1 above and 3 here ensures at least one batch contains # MachineLease entities created for two different MachineTypes. lease_management.drain_excess(max_concurrent=3) self.assertEqual(lease_management.MachineLease.query().count(), 5) for machine_lease in lease_management.MachineLease.query(): self.assertTrue(machine_lease.drained)
def test_lease_errors(self): machine_type = lease_management.MachineType( id='errors', leases=[ lease_management.MachineLease(client_request_id='fake-id-1'), lease_management.MachineLease(client_request_id='fake-id-2'), lease_management.MachineLease(client_request_id='fake-id-3'), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='errors', target_size=2, ) machine_type.put() responses = rpc_to_json( machine_provider.BatchedLeaseResponse(responses=[ machine_provider.LeaseResponse( client_request_id='fake-id-1', error=machine_provider.LeaseRequestError.DEADLINE_EXCEEDED, ), machine_provider.LeaseResponse( client_request_id='fake-id-2', error=machine_provider.LeaseRequestError.REQUEST_ID_REUSE, ), machine_provider.LeaseResponse( client_request_id='fake-id-3', error=machine_provider.LeaseRequestError.TRANSIENT_ERROR, ), ])) lease_management.update_leases(machine_type.key, responses) updated_machine_type = machine_type.key.get() self.assertEqual(len(updated_machine_type.leases), 2) request_ids = sorted(request.client_request_id for request in updated_machine_type.leases) self.assertEqual(request_ids[0], 'fake-id-1') self.assertEqual(request_ids[1], 'fake-id-3')
def test_drain_all(self): key = lease_management.MachineType( enabled=False, target_size=3, ).put() lease_management.MachineLease( id='%s-0' % key.id(), machine_type=key, ).put() lease_management.MachineLease( id='%s-1' % key.id(), machine_type=key, ).put() lease_management.MachineLease( id='%s-2' % key.id(), machine_type=key, ).put() lease_management.drain_excess() self.assertEqual(lease_management.MachineLease.query().count(), 3) for machine_lease in lease_management.MachineLease.query(): self.assertTrue(machine_lease.drained)
def test_missing(self): self.mock(lease_management, 'release', lambda *args, **kwargs: True) machine_lease = lease_management.MachineLease( bot_id='bot-id', client_request_id='req-id', hostname='bot-id', instruction_ts=utils.utcnow(), machine_type=ndb.Key(lease_management.MachineType, 'mt'), ) machine_lease.put() lease_management.check_for_connection(machine_lease) self.failIf(bot_management.get_info_key(machine_lease.bot_id).get()) self.failIf(machine_lease.key.get().client_request_id) self.failIf(machine_lease.key.get().connection_ts)
def test_machine_type_not_found(self): machine_type = lease_management.MachineType( id='not-found', leases=[ lease_management.MachineLease( client_request_id='fake-id', request_hash='fake-hash', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='not-found', target_size=2, ) requests = lease_management.generate_lease_requests( machine_type.key, 'https://example.com') self.failIf(requests)
def test_machine_type_none_pending(self): machine_type = lease_management.MachineType( id='none-pending', leases=[ lease_management.MachineLease( client_request_id='fake-id', hostname='fake-host', request_hash='fake-hash', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_id_base='none-pending', target_size=2, ) requests = lease_management._generate_lease_request_status_updates( machine_type, 'https://example.com') self.failIf(requests)
def test_creates(self): key = lease_management.MachineLease( id='machine-type-1', hostname='hostname', lease_id='lease-id', lease_expiration_ts=utils.utcnow(), machine_type=ndb.Key(lease_management.MachineType, 'machine-type'), ).put() lease_management.ensure_bot_info_exists(key.get()) machine_lease = key.get() bot_info = bot_management.get_info_key(machine_lease.bot_id).get() self.assertEqual(machine_lease.bot_id, machine_lease.hostname) self.assertEqual(bot_info.lease_id, machine_lease.lease_id) self.assertEqual(bot_info.lease_expiration_ts, machine_lease.lease_expiration_ts) self.assertEqual(bot_info.machine_type, machine_lease.machine_type.id()) self.assertEqual(bot_info.machine_lease, machine_lease.key.id())
def test_machine_lease_exists_mismatched_updated(self): def fetch_machine_types(): return { 'machine-type': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type', target_size=1, ), } self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) key = lease_management.MachineType( id='machine-type', early_release_secs=0, lease_duration_secs=1, mp_dimensions=machine_provider.Dimensions(disk_gb=100, ), target_size=1, ).put() key = lease_management.MachineLease( id='%s-0' % key.id(), early_release_secs=1, lease_duration_secs=2, lease_expiration_ts=utils.utcnow(), machine_type=key, mp_dimensions=machine_provider.Dimensions(disk_gb=200, ), ).put() lease_management.ensure_entities_exist() self.assertEqual(lease_management.MachineLease.query().count(), 1) self.assertEqual(key.get().early_release_secs, 0) self.assertEqual(key.get().lease_duration_secs, 1) self.assertEqual(key.get().mp_dimensions.disk_gb, 100)
def test_machine_lease_exists_mismatched_not_updated(self): key = lease_management.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=machine_provider.Dimensions(disk_gb=100, ), target_size=1, ).put() key = lease_management.MachineLease( id='%s-0' % key.id(), early_release_secs=1, lease_duration_secs=2, machine_type=key, mp_dimensions=machine_provider.Dimensions(disk_gb=200, ), ).put() lease_management.ensure_entities_exist() self.assertEqual(lease_management.MachineLease.query().count(), 1) self.assertEqual(key.get().early_release_secs, 1) self.assertEqual(key.get().lease_duration_secs, 2) self.assertEqual(key.get().mp_dimensions.disk_gb, 200)
def test_machine_type_not_enabled(self): machine_type = lease_management.MachineType( id='not-enabled', enabled=False, leases=[ lease_management.MachineLease( client_request_id='fake-id-1', request_hash='fake-hash', ), ], mp_dimensions=machine_provider.Dimensions( os_family=machine_provider.OSFamily.LINUX, ), request_count=1, request_id_base='not-found', target_size=2, ) machine_type.put() requests = lease_management.generate_lease_requests( machine_type.key, 'https://example.com') self.assertEqual(len(requests), 1) self.assertEqual(requests[0].request_id, 'fake-id-1')