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_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_pending_deletion(self):
bot_ids = []
for i in xrange(200):
bot_ids.append('fake-bot-%d' % i)
bot_management.BotInfo(key=bot_management.get_info_key(bot_ids[-1])).put()
machine_type = lease_management.MachineType(
id='few-pending-deletion',
pending_deletion=bot_ids,
request_id_base='few-pending-deletion',
target_size=2,
)
machine_type.put()
lease_management.clean_up_bots()
updated_machine_type = machine_type.key.get()
self.failIf(updated_machine_type.pending_deletion)
for bot_id in bot_ids:
self.failIf(bot_management.get_info_key(bot_id).get())
def test_pending_deletion(self):
bot_ids = []
for i in xrange(200):
bot_ids.append('fake-bot-%d' % i)
bot_management.BotInfo(
key=bot_management.get_info_key(bot_ids[-1])).put()
machine_type = lease_management.MachineType(
id='few-pending-deletion',
pending_deletion=bot_ids,
request_id_base='few-pending-deletion',
target_size=2,
)
machine_type.put()
lease_management.clean_up_bots()
updated_machine_type = machine_type.key.get()
self.failIf(updated_machine_type.pending_deletion)
for bot_id in bot_ids:
self.failIf(bot_management.get_info_key(bot_id).get())
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_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_provider_bots(self):
ndb.get_context().set_cache_policy(lambda _: None)
now = utils.utcnow()
bots = [
bot_management.BotInfo(
key=bot_management.get_info_key('bot1'),
machine_type='machine-type-1',
last_seen_ts=now,
),
bot_management.BotInfo(
key=bot_management.get_info_key('bot2'),
machine_type='machine-type-1',
last_seen_ts=now,
),
bot_management.BotInfo(
key=bot_management.get_info_key('bot3'),
machine_type='machine-type-2',
last_seen_ts=now,
task_id='task',
),
bot_management.BotInfo(
key=bot_management.get_info_key('bot4'),
machine_type='machine-type-3',
last_seen_ts=now,
task_id='task',
),
bot_management.BotInfo(
key=bot_management.get_info_key('bot5'),
machine_type='machine-type-3',
last_seen_ts=now,
),
bot_management.BotInfo(
key=bot_management.get_info_key('bot6'),
machine_type='machine-type-3',
last_seen_ts=now,
task_id='task',
),
]
ndb.put_multi(bots)
obj1 = lease_management.MachineType(id='machine-type-1', target_size=2)
obj1.put()
obj2 = lease_management.MachineType(id='machine-type-2', target_size=1)
obj2.put()
obj3 = lease_management.MachineType(id='machine-type-3', target_size=1)
obj3.put()
lease_management.compute_utilization()
u1 = ndb.Key(lease_management.MachineTypeUtilization,
obj1.key.string_id()).get()
self.assertEqual(u1.busy, 0)
self.assertEqual(u1.idle, 2)
self.failUnless(u1.last_updated_ts)
u2 = ndb.Key(lease_management.MachineTypeUtilization,
obj2.key.string_id()).get()
self.assertEqual(u2.busy, 1)
self.assertEqual(u2.idle, 0)
self.failUnless(u2.last_updated_ts)
u3 = ndb.Key(lease_management.MachineTypeUtilization,
obj3.key.string_id()).get()
self.assertEqual(u3.busy, 2)
self.assertEqual(u3.idle, 1)
self.failUnless(u3.last_updated_ts)