move setting wait_for_fut result to WaitFor await in case it isn't awaited

This commit is contained in:
jesopo 2020-04-30 11:22:47 +01:00
parent 47a4f99c3b
commit bc70afe04b
2 changed files with 17 additions and 11 deletions

View file

@ -14,15 +14,23 @@ class MaybeAwait(Generic[TEvent]):
return coro.__await__()
class WaitFor(object):
def __init__(self, response: IMatchResponse):
self.response = response
self._fut: "Future[Line]" = Future()
def __init__(self,
wait_fut: "Future[WaitFor]",
response: IMatchResponse):
self._wait_fut = wait_fut
self.response = response
self.deferred = False
self._our_fut: "Future[Line]" = Future()
def __await__(self) -> Generator[Any, None, Line]:
return self._fut.__await__()
self._wait_fut.set_result(self)
return self._our_fut.__await__()
async def defer(self):
self.deferred = True
await self
def match(self, server: IServer, line: Line):
return self.response.match(server, line)
def resolve(self, line: Line):
self._fut.set_result(line)
self._our_fut.set_result(line)

View file

@ -258,9 +258,9 @@ class Server(IServer):
if wait_for is not None:
break
async def wait_for(self,
def wait_for(self,
response: Union[IMatchResponse, Set[IMatchResponse]]
) -> Line:
) -> Awaitable[Line]:
response_obj: IMatchResponse
if isinstance(response, set):
response_obj = ResponseOr(*response)
@ -270,10 +270,8 @@ class Server(IServer):
wait_for_fut = self._wait_for_fut
if wait_for_fut is not None:
self._wait_for_fut = None
our_wait_for = WaitFor(response_obj)
wait_for_fut.set_result(our_wait_for)
return await our_wait_for
our_wait_for = WaitFor(wait_for_fut, response_obj)
return our_wait_for
raise Exception()
async def _on_send_line(self, line: Line):