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())