def test_heartbeat_batch(self, heartbeat_batch): self.test_build.lease_expiration_date = self.future_date build2 = model.Build( id=2, bucket='chromium', lease_expiration_date=self.future_date, ) heartbeat_batch.return_value = [ (self.test_build.key.id(), self.test_build, None), (build2.key.id(), None, errors.LeaseExpiredError()) ] req = { 'heartbeats': [{ 'build_id': self.test_build.key.id(), 'lease_key': 42, 'lease_expiration_ts': self.future_ts, }, { 'build_id': build2.key.id(), 'lease_key': 42, 'lease_expiration_ts': self.future_ts, }], } res = self.call_api('heartbeat_batch', req).json_body heartbeat_batch.assert_called_with([{ 'build_id': self.test_build.key.id(), 'lease_key': 42, 'lease_expiration_date': self.future_date, }, { 'build_id': build2.key.id(), 'lease_key': 42, 'lease_expiration_date': self.future_date, }]) result1 = res['results'][0] self.assertEqual(int(result1['build_id']), self.test_build.key.id()) self.assertEqual(result1['lease_expiration_ts'], self.future_ts) result2 = res['results'][1] self.assertEqual(int(result2['build_id']), build2.key.id()) self.assertTrue(result2['error']['reason'] == 'LEASE_EXPIRED')
def test_heartbeat_batch(self, heartbeat_batch): build1 = test_util.build(id=1) build1.lease_expiration_date = self.future_date heartbeat_batch.return_value = [ (1, build1, None), (2, None, errors.LeaseExpiredError()), ] req = { 'heartbeats': [ { 'build_id': '1', 'lease_key': 42, 'lease_expiration_ts': self.future_ts, }, { 'build_id': '2', 'lease_key': 42, 'lease_expiration_ts': self.future_ts, }, ] } res = self.call_api('heartbeat_batch', req).json_body heartbeat_batch.assert_called_with([ { 'build_id': 1, 'lease_key': 42, 'lease_expiration_date': self.future_date, }, { 'build_id': 2, 'lease_key': 42, 'lease_expiration_date': self.future_date, }, ]) result1 = res['results'][0] self.assertEqual(result1['build_id'], '1') self.assertEqual(result1['lease_expiration_ts'], self.future_ts) result2 = res['results'][1] self.assertEqual(result2['build_id'], '2') self.assertTrue(result2['error']['reason'] == 'LEASE_EXPIRED')
def _check_lease(build, lease_key): if lease_key != build.lease_key: raise errors.LeaseExpiredError( 'lease_key for build %s is incorrect. Your lease might be expired.' % build.key.id())