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)