Reimplement Machine Provider integration

appengine/swarming/server/lease_management_test.py

 #!/usr/bin/python
 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Unit tests for lease_management.py."""
 
 import datetime
 import json
 import unittest
@@ skipping 18 matching lines @@
 
   Args:
     rpc_message: A protorpc.message.Message instance.
 
   Returns:
     A string representing a JSON dict.
   """
   return json.loads(protojson.encode_message(rpc_message))
 
 
-class CleanUpBotsTest(test_case.TestCase):
-  """Tests for lease_management.clean_up_bots."""
+class DrainExcessTest(test_case.TestCase):
+  """Tests for lease_management.drain_excess."""
 
   def test_no_machine_types(self):
-    lease_management.clean_up_bots()
+    lease_management.drain_excess()
 
-  def test_none_pending_deletion(self):
-    machine_type = lease_management.MachineType(
-        id='none-pending-deletion',
+    self.failIf(lease_management.MachineLease.query().count())
+
+  def test_nothing_to_drain(self):
+    key = lease_management.MachineType(
+        target_size=1,
+    ).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.failIf(key.get().drained)
+
+  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_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_drain_batched(self):
+    key = lease_management.MachineType(
+        enabled=False,
         target_size=2,
-    )
-    machine_type.put()
+    ).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()
 
-    lease_management.clean_up_bots()
-    updated_machine_type = machine_type.key.get()
-    self.failIf(updated_machine_type.pending_deletion)
+    # 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)
 
-  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,
-        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())
+    self.assertEqual(lease_management.MachineLease.query().count(), 5)
+    for machine_lease in lease_management.MachineLease.query():
+      self.assertTrue(machine_lease.drained)
 
 
-class ClearBotsPendingDeletionTest(test_case.TestCase):
-  """Tests for lease_management._clear_bots_pending_deletion."""
 
-  def test_machine_type_not_found(self):
-    machine_type = lease_management.MachineType(
-        id='not-found',
-        pending_deletion=['fake-bot-1', 'fake-bot-2'],
-    )
+class EnsureEntitiesExistTest(test_case.TestCase):
+  """Tests for lease_management.ensure_entities_exist."""
 
-    lease_management._clear_bots_pending_deletion(
-        machine_type.key, ['fake-bot-1', 'fake-bot-2'])
-    self.assertEqual(len(machine_type.pending_deletion), 2)
+  def test_no_machine_types(self):
+    lease_management.ensure_entities_exist()
 
-  def test_machine_type_none_pending_deletion(self):
-    machine_type = lease_management.MachineType(
-        id='none-pending-deletion',
+    self.failIf(lease_management.MachineLease.query().count())
+
+  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_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)
+
+  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_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_three_machine_types_multiple_batches(self):
+    lease_management.MachineType(
+        id='a',
         target_size=2,
-    )
-    machine_type.put()
+    ).put()
+    lease_management.MachineType(
+        id='b',
+        target_size=2,
+    ).put()
+    lease_management.MachineType(
+        id='c',
+        target_size=1,
+    ).put()
 
-    lease_management._clear_bots_pending_deletion(machine_type.key, [])
-    updated_machine_type = machine_type.key.get()
-    self.failIf(updated_machine_type.pending_deletion)
+    # 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)
 
-  def test_machine_type_pending_deletion(self):
-    machine_type = lease_management.MachineType(
-        id='pending-deletion',
-        pending_deletion=['fake-bot-1', 'fake-bot-2'],
-        target_size=2,
-    )
-    machine_type.put()
-
-    lease_management._clear_bots_pending_deletion(
-        machine_type.key, ['fake-bot-1'])
-    updated_machine_type = machine_type.key.get()
-    self.assertEqual(len(updated_machine_type.pending_deletion), 1)
-    self.assertEqual(updated_machine_type.pending_deletion[0], 'fake-bot-2')
-
-
-class GenerateLeaseRequestsTest(test_case.TestCase):
-  """Tests for lease_management.generate_lease_requests."""
-
-  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,
-        ),
-        target_size=2,
-    )
-
-    requests = lease_management.generate_lease_requests(
-        machine_type.key, 'example', 'https://example.com')
-    self.failIf(requests)
-
-  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,
-        target_size=2,
-    )
-    machine_type.put()
-
-    requests = lease_management.generate_lease_requests(
-        machine_type.key, 'example', 'https://example.com')
-    self.assertEqual(len(requests), 1)
-    self.assertEqual(requests[0].request_id, 'fake-id-1')
-
-  def test_need_one(self):
-    machine_type = lease_management.MachineType(
-      id='need-one',
-      mp_dimensions=machine_provider.Dimensions(
-          os_family=machine_provider.OSFamily.LINUX,
-      ),
-      target_size=1,
-    )
-    machine_type.put()
-
-    requests = lease_management.generate_lease_requests(
-        machine_type.key, 'example', 'https://example.com')
-    updated_machine_type = machine_type.key.get()
-    self.assertEqual(len(requests), 1)
-    self.assertEqual(len(updated_machine_type.leases), 1)
-    self.assertEqual(
-        updated_machine_type.leases[0].client_request_id, 'need-one-1')
-    self.assertEqual(updated_machine_type.request_count, 1)
-
-
-class CleanUpExpiredLeasesTest(test_case.TestCase):
-  """Tests for lease_management._clean_up_expired_leases."""
-
-  def test_machine_type_no_leases(self):
-    machine_type = lease_management.MachineType(
-        id='no-leases',
-        leases=[],
-        mp_dimensions=machine_provider.Dimensions(
-            os_family=machine_provider.OSFamily.LINUX,
-        ),
-        target_size=2,
-    )
-
-    expired  = lease_management._clean_up_expired_leases(machine_type)
-    self.failIf(expired)
-    self.failIf(machine_type.leases)
-    self.failIf(machine_type.pending_deletion)
-
-  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,
-        ),
-        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_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',
-        ],
-        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')
-
-
-class GenerateLeaseRequestStatusUpdatesTest(test_case.TestCase):
-  """Tests for lease_management._generate_lease_request_status_updates."""
-
-  def test_machine_type_no_leases(self):
-    machine_type = lease_management.MachineType(
-        id='no-leases',
-        leases=[],
-        mp_dimensions=machine_provider.Dimensions(
-            os_family=machine_provider.OSFamily.LINUX,
-        ),
-        target_size=2,
-    )
-
-    requests = lease_management._generate_lease_request_status_updates(
-        machine_type, 'example', '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,
-        ),
-        target_size=2,
-    )
-
-    requests = lease_management._generate_lease_request_status_updates(
-        machine_type, 'example', 'https://example.com')
-    self.failIf(requests)
-
-  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,
-        ),
-        target_size=2,
-    )
-
-    requests = lease_management._generate_lease_request_status_updates(
-        machine_type, 'example', 'https://example.com')
-    self.assertEqual(len(requests), 1)
-    self.assertEqual(requests[0].request_id, 'fake-id-3')
-
-
-class GenerateLeaseRequestsForNewMachinesTest(test_case.TestCase):
-  """Tests for lease_management._generate_lease_requests_for_new_machines."""
-
-  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,
-        ),
-        target_size=2,
-    )
-
-    requests = lease_management._generate_lease_requests_for_new_machines(
-        machine_type, 'example', 'https://example.com')
-    self.failIf(requests)
-
-  def test_need_one(self):
-    machine_type = lease_management.MachineType(
-      id='need-one',
-      mp_dimensions=machine_provider.Dimensions(
-          os_family=machine_provider.OSFamily.LINUX,
-      ),
-      target_size=1,
-    )
-
-    requests = lease_management._generate_lease_requests_for_new_machines(
-        machine_type, 'example', 'https://example.com')
-    self.assertEqual(len(requests), 1)
-
-  def test_ensure_correct_request_ids(self):
-    machine_type = lease_management.MachineType(
-      id='ensure-correct-request-ids',
-      mp_dimensions=machine_provider.Dimensions(
-          os_family=machine_provider.OSFamily.LINUX,
-      ),
-      request_count=2,
-      target_size=2,
-    )
-
-    requests = lease_management._generate_lease_requests_for_new_machines(
-        machine_type, 'example', 'https://example.com')
-    self.assertEqual(len(requests), 2)
-    request_ids = sorted(request.request_id for request in requests)
-    self.assertEqual(request_ids[0], 'ensure-correct-request-ids-3')
-    self.assertEqual(request_ids[1], 'ensure-correct-request-ids-4')
-
-
-class UpdateLeasesTest(test_case.TestCase):
-  """Tests for lease_management.update_leases."""
-
-  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,
-      ),
-      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['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_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,
-      ),
-      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['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_lease_untriaged(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,
-      ),
-      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.UNTRIAGED,
-        ),
-    ]))
-
-    lease_management.update_leases(machine_type.key, responses['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-1')
-    self.failIf(updated_machine_type.leases[0].hostname)
-    self.failIf(updated_machine_type.leases[0].lease_expiration_ts)
-    self.assertEqual(updated_machine_type.leases[0].request_hash, 'fake-hash-1')
-
-  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,
-      ),
-      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['responses'])
-    updated_machine_type = machine_type.key.get()
-    self.failIf(updated_machine_type.leases)
-
-  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,
-      ),
-      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['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')
+    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'))
 
 
 if __name__ == '__main__':
   unittest.main()