move setting wait_for_fut result to WaitFor await in case it isn't awaited
This commit is contained in:
parent
47a4f99c3b
commit
bc70afe04b
2 changed files with 17 additions and 11 deletions
|
@ -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)
|
||||
|
|
|
@ -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):
|
||||
|
|
Loading…
Reference in a new issue