async def on_websocket(self, req, ws): self.times_called += 1 await ws.accept() while True: try: if direction == 'send': await ws.send_data(sample_data) if explicit_close_server: await ws.close(4099) else: self.data = await ws.receive_data() if explicit_close_server: # NOTE(kgriffs): We call ws.receive_data() again here in order # to test coverage of the logic that handles the case # of a closed connection while waiting on more data. recv_task = falcon.create_task(ws.receive_data()) await asyncio.sleep( 0 ) # Ensure recv_task() has a chance to get ahead await asyncio.wait([recv_task, ws.close(4099)]) self.data_received.set() except falcon.WebSocketDisconnected as ex: self.ws_ready = ws.ready self.ws_client_close_code = ex.code self.ws_state = ws._state self.ws_closed = ws.closed break
async def on_websocket(self, req, ws, p1, p2, injected): # NOTE(kgriffs): Normally the receiver task is not started # until the websocket is started. But here we start it # first in order to simulate a potential race condition # that ASGIWebSocketSimulator._emit() guards against, # in case it ever arises due to the way the target ASGI # app may be implemented. ws._buffered_receiver.start() await asyncio.sleep(0) if ws.unaccepted: await ws.accept() try: await ws.accept() except falcon.OperationNotAllowed: self.caught_operation_not_allowed = True if ws.ready: await ws.send_text( f'{p1}:{p2}:{req.context.message}:{injected}') messages = deque() sink_task = falcon.create_task(self._sink(ws, messages)) while not sink_task.done(): if not messages: await asyncio.sleep(0) continue try: await ws.send_text(messages.popleft()) except falcon.WebSocketDisconnected: break sink_task.cancel() try: await sink_task except asyncio.CancelledError: pass
def start(self): if not self._pump_task: self._pump_task = falcon.create_task(self._pump())