def test_one_enabled_machine_type(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', target_size=1, ).put() lease_management.ensure_entities_exist() self.assertEqual(key.get().early_release_secs, 0) self.assertEqual(key.get().lease_duration_secs, 1) self.assertEqual(key.get().target_size, 1) self.assertEqual(lease_management.MachineLease.query().count(), 1)
def test_enable_and_update_machine_type(self): def fetch_machine_types(): return { 'machine-type': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=2, 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, enabled=False, lease_duration_secs=1, mp_dimensions=machine_provider.Dimensions(disk_gb=100, ), target_size=1, ).put() lease_management.ensure_entities_exist() self.failUnless(key.get().enabled) self.assertEqual(key.get().lease_duration_secs, 2)
def test_no_enabled_machine_types(self): lease_management.MachineType( enabled=False, target_size=3, ).put() lease_management.ensure_entities_exist() self.failIf(lease_management.MachineLease.query().count())
def test_two_enabled_machine_types(self): lease_management.MachineType( id='a', target_size=1, ).put() lease_management.MachineType( id='b', target_size=1, ).put() lease_management.ensure_entities_exist() self.assertEqual(lease_management.MachineLease.query().count(), 2) self.failUnless(lease_management.MachineLease.get_by_id('a-0')) self.failUnless(lease_management.MachineLease.get_by_id('b-0'))
def test_daily_schedule_resize_to_zero(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, schedule=bots_pb2.Schedule(daily=[ bots_pb2.DailySchedule( start='0:00', end='1:00', days_of_the_week=xrange(7), target_size=0, ) ], ), ), } self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) self.mock( lease_management.utils, 'utcnow', lambda: datetime.datetime(1969, 1, 1, 0, 30), ) 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() lease_management.ensure_entities_exist() self.failIf(lease_management.MachineLease.query().count()) self.failIf(key.get().target_size)
def test_two_enabled_machine_types(self): def fetch_machine_types(): return { 'machine-type-a': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type-a', target_size=1, ), 'machine-type-b': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type-b', target_size=1, ), } self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) lease_management.MachineType( id='machine-type-a', target_size=1, ).put() lease_management.MachineType( id='machine-type-b', target_size=1, ).put() lease_management.ensure_entities_exist() self.assertEqual(lease_management.MachineLease.query().count(), 2) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-a-0')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-b-0'))
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_disable_machine_type(self): def fetch_machine_types(): return {} self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) key = lease_management.MachineType( id='machine-type', early_release_secs=0, enabled=True, lease_duration_secs=1, mp_dimensions=machine_provider.Dimensions(disk_gb=100, ), target_size=1, ).put() lease_management.ensure_entities_exist() self.failIf(key.get().enabled)
def test_one_machine_type_multiple_batches(self): lease_management.MachineType( id='machine-type', target_size=5, ).put() # Choice of 3 here and 5 above ensures MachineLeases are created in two # batches of differing sizes. lease_management.ensure_entities_exist(max_concurrent=3) self.assertEqual(lease_management.MachineLease.query().count(), 5) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-0')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-1')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-2')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-3')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-4'))
def test_one_machine_type_multiple_batches(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=5, ), } self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) lease_management.MachineType( id='machine-type', target_size=5, ).put() # Choice of 3 here and 5 above ensures MachineLeases are created in two # batches of differing sizes. lease_management.ensure_entities_exist(max_concurrent=3) self.assertEqual(lease_management.MachineLease.query().count(), 5) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-0')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-1')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-2')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-3')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-4'))
def test_three_machine_types_multiple_batches(self): lease_management.MachineType( id='a', target_size=2, ).put() lease_management.MachineType( id='b', target_size=2, ).put() lease_management.MachineType( id='c', target_size=1, ).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.ensure_entities_exist(max_concurrent=3) self.assertEqual(lease_management.MachineLease.query().count(), 5) self.failUnless(lease_management.MachineLease.get_by_id('a-0')) self.failUnless(lease_management.MachineLease.get_by_id('a-1')) self.failUnless(lease_management.MachineLease.get_by_id('b-0')) self.failUnless(lease_management.MachineLease.get_by_id('b-1')) self.failUnless(lease_management.MachineLease.get_by_id('c-0'))
def test_three_machine_types_multiple_batches(self): def fetch_machine_types(): return { 'machine-type-a': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type-a', target_size=2, ), 'machine-type-b': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type-b', target_size=2, ), 'machine-type-c': bots_pb2.MachineType( early_release_secs=0, lease_duration_secs=1, mp_dimensions=['disk_gb:100'], name='machine-type-c', target_size=1, ), } self.mock( lease_management.bot_groups_config, 'fetch_machine_types', fetch_machine_types, ) lease_management.MachineType( id='machine-type-a', target_size=2, ).put() lease_management.MachineType( id='machine-type-b', target_size=2, ).put() lease_management.MachineType( id='machine-type-c', target_size=1, ).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.ensure_entities_exist(max_concurrent=3) self.assertEqual(lease_management.MachineLease.query().count(), 5) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-a-0')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-a-1')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-b-0')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-b-1')) self.failUnless( lease_management.MachineLease.get_by_id('machine-type-c-0'))
def test_no_machine_types(self): lease_management.ensure_entities_exist() self.failIf(lease_management.MachineLease.query().count())
def test_one_enabled_machine_type(self): lease_management.MachineType(target_size=1, ).put() lease_management.ensure_entities_exist() self.assertEqual(lease_management.MachineLease.query().count(), 1)