def test_poll_task_exception(self, mCreateSession):
iSession = MockSession()
mCreateSession.return_value = (iSession, '123456')
client = iSession.client('stepfunctions')
activity = ActivityMixin()
self.assertEqual(activity.arn, None)
with self.assertRaises(Exception):
activity.poll_task(worker='worker')
calls = []
self.assertEqual(client.mock_calls, calls)
def test_poll_task_no_work(self, mCreateSession):
iSession = MockSession()
mCreateSession.return_value = (iSession, '123456')
client = iSession.client('stepfunctions')
client.get_activity_task.return_value = {'taskToken': ''}
activity = ActivityMixin()
activity.arn = 'XXX'
token, input_ = activity.poll_task('worker')
self.assertEqual(token, None)
self.assertEqual(input_, None)
calls = [
mock.call.get_activity_task(activityArn='XXX', workerName='worker')
]
self.assertEqual(client.mock_calls, calls)