test_activities.py 文件源码

def test_run_generator(self, mCreateSession):
        iSession = MockSession()
        mCreateSession.return_value = (iSession, '123456')
        client = iSession.client('stepfunctions')

        def target(input_):
            yield
            yield
            return

        # Just make sure the target is actually a generator
        self.assertEqual(type(target(None)), types.GeneratorType)

        task = TaskMixin(process = target)
        task.handle_task('token', None)

        self.assertEqual(task.token, None)
        call = mock.call.send_task_success(taskToken = 'token',
                                           output = 'null')
        call_ = mock.call.send_task_heartbeat(taskToken = 'token')
        calls = [call_, call_, call]
        self.assertEqual(client.mock_calls, calls)