Verified Commit 1a67a286 authored by Dom Sekotill's avatar Dom Sekotill
Browse files

Add a "started" class context to MockMessageStream 

This removes some of the boilerplate associated with MockMessageStream
parent eed4ac66

tests/mock_stream.py

+12 −0
Original line number Diff line number Diff line
@@ -19,6 +19,7 @@ from typing_extensions import Self 
from kilter.protocol import *

from kilter.protocol.buffer import SimpleBuffer

from kilter.service import ResponseMessage

from kilter.service import Runner



P = typing.ParamSpec("P")

SendT = typing.TypeVar("SendT")
@@ -75,6 +76,17 @@ class MockMessageStream: 
		await self._stream.__aexit__(et, ex, tb)

		await self.peer_stream.__aexit__(et, ex, tb)



	@classmethod

	@asynccontextmanager

	async def started(cls, runner: Runner) -> AsyncIterator[Self]:

		"""

		Return a context manager that yields a prepared stream mock connected to a runner

		"""

		async with anyio.create_task_group() as tg, cls() as stream_mock:

			tg.start_soon(runner, stream_mock.peer_stream)

			await anyio.wait_all_tasks_blocked()

			yield stream_mock



	async def abort(self) -> None:

		"""

		Send Abort and close the stream