async def test_client_open_url(echo_server):
url = 'ws://{}:{}{}/path'.format(HOST, echo_server.port, RESOURCE)
async with open_websocket_url(url) as conn:
assert conn.path == RESOURCE + '/path'
url = 'ws://{}:{}{}?foo=bar'.format(HOST, echo_server.port, RESOURCE)
async with open_websocket_url(url) as conn:
assert conn.path == RESOURCE + '?foo=bar'
async def send_iter(addr, iterable):
try:
async with open_websocket_url(addr) as ws:
for message in iterable:
await ws.send_message(message)
except OSError as ose:
print('Connection attempt failed: %s' % ose, file=stderr)
async def stream_ticks(tickers: str,
output_fname: str,
print_msg: bool = False,
batch: bool = True):
# connect to ws
async with open_websocket_url(
'wss://alpaca.socket.polygon.io/stocks') as ws:
# authenticate
await ws.send_message('{"action":"auth", "params":"' +
POLYGON_API_KEY + '"}')
# subscribe to symbols
await ws.send_message(
f'{{"action": "subscribe", "params": "{tickers}"}}')
# listen for events
while True:
with open(output_fname, 'a') as out_file:
message = await ws.get_message()
if print_msg:
print(message)
if batch:
out_file.write(str(time_ns()) + '||' + message + '\n')
else:
ticks = json.loads(message)
for tick in ticks:
if tick['ev'] not in ['T', 'Q']:
continue
tick['a'] = time_ns()
out_file.write(json.dumps(tick) + '\n')
return True
async def send_bus_updates(server_address: str,
receive_channel: trio.MemoryReceiveChannel) -> None:
"""Consume updates from trio channel and send them to a server via websockets."""
async with open_websocket_url(server_address) as ws:
logger.debug(f'Established connection with {server_address}')
async for message in receive_channel:
await ws.send_message(json.dumps(message, ensure_ascii=False))
async def open_jsonrpc_ws(url: str) -> typing.AsyncIterator[JsonRpcConnection]:
""" Open a JSON-RPC connection using WebSocket transport. """
async with trio_websocket.open_websocket_url(url) as ws:
async with trio.open_nursery() as nursery:
transport = WebSocketTransport(ws)
yield jsonrpc_client(transport, nursery)
nursery.cancel_scope.cancel()
async def send_one(addr, message):
try:
async with open_websocket_url(addr) as ws:
await ws.send_message(message)
message = await ws.get_message()
print('Received message: %s' % message)
except OSError as ose:
print('Connection attempt failed: %s' % ose, file=stderr)
async def make_request(url, message):
try:
async with open_websocket_url(url) as ws:
await ws.send_message(message)
message = await ws.get_message()
return message
except OSError as ose:
logging.error('Connection attempt failed: %s', ose)
async def main():
try:
async with open_websocket_url('ws://127.0.0.1:8060/ws') as ws:
await ws.send_message('hello world!')
message = await ws.get_message()
print('Received message: %s' % message)
except OSError as ose:
print('Connection attempt failed: %s' % ose)
async def main():
try:
async with open_websocket_url('ws://localhost:8000/') as ws:
async with trio.open_nursery() as nursery:
nursery.start_soon(web_socket_to_stdout, ws)
nursery.start_soon(stdin_to_web_socket, ws)
except OSError as ose:
print('Connection attempt failed: %s' % ose, file=stderr)
async def _websocket_connect(connection):
if connection is None:
raise ZoomError("Connection rejected by Zoom")
websocket_url = str(connection.url).replace("https", "wss")
logger.debug("WebSocket connection url: {websocket_url}",
websocket_url=websocket_url)
return open_websocket_url(websocket_url)
async def open_ws_wrapper(addr, on_connect, on_error=_default_on_error):
try:
async with open_websocket_url(addr) as ws:
on_connect(ws)
except OSError as ose:
on_error(ose)
except ConnectionClosed as e:
on_error(e)
async def run(self):
try:
async with open_websocket_url("ws://localhost:1234") as ws:
self._ws = ws
while True:
msg = await ws.get_message()
await self.process_message(msg)
except OSError:
print("Cannot connect to minecraft")
async def several_reqs(url, msgs):
try:
async with open_websocket_url(url) as ws:
for m in msgs:
await ws.send_message(json.dumps(m))
message = await ws.get_message()
return message
except OSError as ose:
logging.error('Connection attempt failed: %s', ose)
async def main():
try:
async with open_websocket_url('ws://127.0.0.1:8000') as ws:
await ws.send_message('{"test": 1}')
message = await ws.get_message()
logging.info('Received message: %s', message)
except OSError as ose:
logging.error('Connection attempt failed: %s', ose)
async def run_case(url, case):
url = URL(url).with_path('/runCase').with_query(case=case, agent=AGENT)
try:
async with open_websocket_url(url) as conn:
while True:
data = await conn.get_message()
await conn.send_message(data)
except ConnectionClosed:
pass
async def run_case(url, case):
url = f'{url}/runCase?case={case}&agent={AGENT}'
try:
async with open_websocket_url(url, max_message_size=MAX_MESSAGE_SIZE) as conn:
while True:
data = await conn.get_message()
await conn.send_message(data)
except ConnectionClosed:
pass
async def main(url: str):
async with open_websocket_url(url) as ws:
async with await trio.open_file("messages") as f:
await f.seek(0, 2) # go to end of file
while True:
line = await f.readline()
if not line:
continue
await ws.send_message(line.strip())
message = await ws.get_message()
print(message)
async def main():
try:
async with open_websocket_url('ws://127.0.0.1:8080/ws') as ws:
logging.info(f'Sending different bad messages to server')
for desc, message in _BAD_MESSAGES.items():
logging.info(f'Sending message "{desc}"')
await ws.send_message(json.dumps(message, ensure_ascii=False))
message = await ws.get_message()
logging.info(f'Received answer: {message}')
except OSError as ose:
logging.error(f'Connection attempt failed: {ose }')
def test_client_open_url_options(open_websocket_mock):
"""open_websocket_url() must pass its options on to open_websocket()"""
port = 1234
url = 'ws://{}:{}{}'.format(HOST, port, RESOURCE)
options = {
'subprotocols': ['chat'],
'extra_headers': [(b'X-Test-Header', b'My test header')],
'message_queue_size': 9,
'max_message_size': 333,
'connect_timeout': 36,
'disconnect_timeout': 37,
}
open_websocket_url(url, **options)
_, call_args, call_kwargs = open_websocket_mock.mock_calls[0]
assert call_args == (HOST, port, RESOURCE)
assert not call_kwargs.pop('use_ssl')
assert call_kwargs == options
open_websocket_url(url.replace('ws:', 'wss:'))
_, call_args, call_kwargs = open_websocket_mock.mock_calls[1]
assert call_kwargs['use_ssl']
async def open_cdp_connection(url) -> typing.AsyncIterator[CdpConnection]:
'''
This async context manager opens a connection to the browser specified by
``url`` before entering the block, then closes the connection when the block
exits.
'''
async with open_websocket_url(url, max_message_size=2**24) as ws:
async with trio.open_nursery() as nursery:
cdp_conn = CdpConnection(ws)
async with cdp_conn:
nursery.start_soon(cdp_conn._reader_task)
yield cdp_conn
async def send_updates(server_address, receive_channel, websocket_number,
refresh_timeout):
async with AsyncExitStack() as stack:
sockets = [
await stack.enter_async_context(open_websocket_url(server_address))
for _ in range(websocket_number)
]
async for value in receive_channel:
try:
await random.choice(sockets).send_message(value)
await trio.sleep(refresh_timeout)
except OSError as ose:
logging.error(f'Connection attempt failed: {ose}')
async def send_m_messages(m):
global websockets_connected
global failed_connections_counter
try:
async with open_websocket_url(url) as ws:
websockets_connected += 1
for _ in range(m):
await send_message_get_response(ws)
except OSError as ose:
failed_connections_counter += 1
logging.error(f'Connection attempt failed {ose}')
finally:
websockets_connected -= 1
async def communicate(addr, iterable_fn):
try:
async with open_websocket_url(addr) as ws:
async for message in iterable_fn(message_gen(ws)):
try:
await ws.send_message(message)
except ConnectionClosed:
print('ConnectionClosed while sending')
return
except OSError as ose:
print('Connection attempt failed: %s' % ose, file=stderr)
async def send_iter(addr, iterable):
try:
async with open_websocket_url(addr) as ws:
try:
async for message in iterable(ws):
if message is not None:
await ws.send_message(message)
except StopIteration:
print("Closing connection")
return
except OSError as ose:
print('Connection attempt failed: %s' % ose, file=stderr)
except ConnectionClosed as e:
print(f'Connection to {addr} closed', e)
async def main():
url = "ws://127.0.0.1:8080"
async with open_websocket_url(url) as ws:
while True:
try:
error_message = random.choice(ERROR_MESSAGES)
await ws.send_message(error_message)
print(f'Sent message "{error_message}" to {url}')
message = await ws.get_message()
print(f"Received message: {message} \n")
except (OSError, ConnectionClosed) as e:
print(f"Connection failed: {e}")
break
await trio.sleep(1)
async def get_worlds():
""" Retrieve a list of worlds from the server
"""
async with open_websocket_url("ws://localhost:1234") as ws:
s = {
"verb": "get",
"type": "worlds",
"location": {
"x": 0,
"y": 0,
"z": 0
}
}
await ws.send_message(json.dumps(s))
return await ws.get_message()
async def connect(self) -> None:
await self.login()
await self.get_world_list()
headers = [('Cookie', f"{AUTH_COOKIE}={self.cookiejar[AUTH_COOKIE]}")]
async with trio_websocket.open_websocket_url(
self.ws_url, extra_headers=headers) as websocket:
self.ws = websocket
self.log('@ Connected')
self.connected = True
await self._callback('on_connected')
async with trio.open_nursery() as nursery:
self.nursery = nursery
nursery.start_soon(self.reader, websocket)
self.connected = False
self.log('@ Disconnected')
await self._callback('on_disconnected')
async def main(log, host, port, imitator_type):
if not log:
logger = logging.getLogger()
logger.disabled = True
url = f'ws://{host}:{port}'
async with open_websocket_url(url) as ws:
while True:
try:
error_message = random.choice(
ERROR_MESSAGES.get(imitator_type))
await ws.send_message(error_message)
logging.info(f'Sent message "{error_message}" to {url}')
message = await ws.get_message()
logging.info(f'Received message: {message}')
except (OSError, ConnectionClosed) as e:
logging.error(f'Connection failed: {e}')
break
await trio.sleep(1)
async def run_bus(send_channel, server_address, route, bus_id):
first_run = True
while True:
if first_run:
async for serialized_bus_info in get_serialized_bus_info(route,
bus_id):
try:
async with open_websocket_url(server_address) as ws:
async with send_channel:
if 'busId' in serialized_bus_info: # good connect
helpers.conn_attempt = 0
await ws.send_message(serialized_bus_info)
await trio.sleep(settings['refresh_timeout'])
first_run = False
except OSError as e:
first_run = True
logger.error(f'Connection failed: {e=}')
except HandshakeError:
pass
async def main():
try:
async with open_websocket_url('ws://127.0.0.1:8000/ws') as ws:
logging.info(f'Sending different bad messages to server')
for desc, message in _BAD_MESSAGES.items():
logging.info(f'Sending message "{desc}"')
await ws.send_message(json.dumps(message, ensure_ascii=False))
# note - server sends not only errors but buses positions too.
# so if we just use ws.get_message() we will get just 'msgType': 'Buses' message
with trio.move_on_after(2) as cancel_scope:
message = ''
while 'Errors' not in message:
message = await ws.get_message()
logging.info(f'Received answer: {message}')
if cancel_scope.cancelled_caught:
logging.error(
'Server should returned error message but it didn\'t!')
except OSError as ose:
logging.error(f'Connection attempt failed: {ose}')
async def test_api_local_websockets_stream(local_daemon_app, empty_vault, local_api_client):
app = local_daemon_app
client = local_api_client
c = await client.post('/v1/vault/',
data=json.dumps({ 'folder': empty_vault.folder }))
assert c['resource_uri'] != '/v1/vault/None/'
assert len(c['resource_uri']) > 20
vault_uri = c['resource_uri']
assert len(app.vaults) == 1 # one vault
while app.vaults[0].state in (VaultState.UNINITIALIZED, VaultState.SYNCING):
await trio.sleep(0.2)
assert len(app.vaults) == 1 # one vault
c = await client.get(vault_uri)
assert c['state'] == 'ready'
ws_url = "ws://127.0.0.1:28081" + vault_uri + "historystream/"
async with open_websocket_url(ws_url) as ws:
assert not ws.closed
assert ws.is_client
with trio.move_on_after(0.5):
message = await ws.get_message()
assert False, "should not receive a message"
patch_data = json.dumps({
'metadata': dict(c['metadata'], name='newname')
})
c = await client.put(vault_uri, data=patch_data)
msg = None
with trio.move_on_after(0.5):
msg = await ws.get_message()
assert msg is not None
rev = json.loads(msg)
assert rev['verified'] == True
assert rev['operation'] == "OP_SET_METADATA"