Function test_engine

bbot/test/test_step_1/test_engine.py:5–146

()

@pytest.mark.asyncio
async def test_engine():
    from bbot.core.engine import EngineClient, EngineServer

    counter = 0
    yield_cancelled = False
    yield_errored = False
    return_started = False
    return_finished = False
    return_cancelled = False
    return_errored = False

    class TestEngineServer(EngineServer):
        CMDS = {
            0: "return_thing",
            1: "yield_stuff",
        }

        async def return_thing(self, n):
            nonlocal return_started
            nonlocal return_finished
            nonlocal return_cancelled
            nonlocal return_errored
            try:
                return_started = True
                await asyncio.sleep(n)
                return_finished = True
                return f"thing{n}"
            except asyncio.CancelledError:
                return_cancelled = True
                raise
            except Exception:
                return_errored = True
                raise

        async def yield_stuff(self, n):
            nonlocal counter
            nonlocal yield_cancelled
            nonlocal yield_errored
            try:
                for i in range(n):
                    yield f"thing{i}"
                    counter += 1
                    await asyncio.sleep(0.1)
            except asyncio.CancelledError:
                yield_cancelled = True
                raise
            except Exception:
                yield_errored = True
                raise

    class TestEngineClient(EngineClient):
        SERVER_CLASS = TestEngineServer

        async def return_thing(self, n):
            return await self.run_and_return("return_thing", n)

        async def yield_stuff(self, n):
            async for _ in self.run_and_yield("yield_stuff", n):