def test_lease_reach_latest(self):
"""Test leasing a task and reaching the maximum lease time."""
message = mock.MagicMock()
message.attributes = {
'command': 'cmd',
'argument': 'arg',
'job': 'job',
'eta': time.time(),
}
task = tasks.PubSubTask(message)
event = mock.MagicMock()
event.wait.side_effect = [False, False, False]
self.mock.time.side_effect = [
0, 0, 0, 600, tasks.TASK_LEASE_SECONDS - 30, tasks.TASK_LEASE_SECONDS
]
with task.lease(_event=event) as thread:
self.assertEqual('21600', os.environ['TASK_LEASE_SECONDS'])
thread.join()
self.assertEqual(3, message.modify_ack_deadline.call_count)
# Last extension should be until TASK_LEASE_SECONDS.
message.modify_ack_deadline.assert_has_calls([
mock.call(600),
mock.call(600),
mock.call(30),
])
self.assertEqual(1, message.ack.call_count)
self.assertFalse(thread.is_alive())
def test_lease_exception(self):
"""Test lease with an exception during the task."""
message = mock.MagicMock()
message.attributes = {
'command': 'cmd',
'argument': 'arg',
'job': 'job',
'eta': time.time(),
}
task = tasks.PubSubTask(message)
with self.assertRaises(Exception):
with task.lease() as thread:
raise Exception
self.assertFalse(thread.is_alive())
self.assertEqual(0, message.ack.call_count)
def test_lease_finish_before_latest(self):
"""Test leasing a task and finishing before the latest time."""
message = mock.MagicMock()
message.attributes = {
'command': 'cmd',
'argument': 'arg',
'job': 'job',
'eta': time.time(),
}
task = tasks.PubSubTask(message)
self.mock.time.side_effect = [0, 0, 0, 600, 1200]
with task.lease() as thread:
self.assertEqual('21600', os.environ['TASK_LEASE_SECONDS'])
self.assertEqual(1, message.modify_ack_deadline.call_count)
message.modify_ack_deadline.assert_has_calls([
mock.call(600),
])
self.assertFalse(thread.is_alive())
self.assertEqual(1, message.ack.call_count)
