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())