zeek/auxil/broker/tests/btest/web-socket/gh-357-regression.py
Patrick Kelley 8fd444092b initial
2025-05-07 15:35:15 -04:00

137 lines
3.6 KiB
Python

# @TEST-GROUP: web-socket
#
# @TEST-PORT: BROKER_WEB_SOCKET_PORT
#
# @TEST-COPY-FILE: ${API_ROOT}/v1-broker-out.json
#
# @TEST-EXEC: btest-bg-run node "broker-node --config-file=../node.cfg"
# @TEST-EXEC: btest-bg-run recv "python3 ../recv.py >recv.out"
# @TEST-EXEC: $SCRIPTS/wait-for-file recv/ready 15 || (btest-bg-wait -k 1 && false)
#
# @TEST-EXEC: btest-bg-run send "python3 ../send.py"
#
# @TEST-EXEC: $SCRIPTS/wait-for-file recv/done 30 || (btest-bg-wait -k 1 && false)
# @TEST-EXEC: btest-diff recv/recv.out
#
# @TEST-EXEC: btest-bg-wait -k 1
@TEST-START-FILE node.cfg
caf {
logger {
file {
path = 'broker.log'
verbosity = 'debug'
excluded-components = []
}
}
}
broker {
disable-ssl = true
}
topics = ["/test"]
verbose = true
@TEST-END-FILE
@TEST-START-FILE recv.py
import asyncio, websockets, os, time, json, sys
from jsonschema import validate
ws_port = os.environ['BROKER_WEB_SOCKET_PORT'].split('/')[0]
ws_url = f'ws://localhost:{ws_port}/v1/messages/json'
schema = json.load(open('../v1-broker-out.json'))
async def do_run():
# Try up to 30 times.
connected = False
for i in range(30):
try:
ws = await websockets.connect(ws_url)
connected = True
# send filter and wait for ack
await ws.send('["/test"]')
ack_json = await ws.recv()
ack = json.loads(ack_json)
try:
validate(ack, schema)
except Exception as err:
print(f'received invalid ack: {err}')
sys.exit(1)
# tell btest to start the sender now
with open('ready', 'w') as f:
f.write('ready')
# the message must be valid JSON
msg_json = await ws.recv()
msg = json.loads(msg_json)
try:
validate(msg, schema)
except Exception as err:
print(f'received invalid data message: {err}')
sys.exit(1)
validate(msg, schema)
# pretty-print to stdout (redirected to recv.out)
print(json.dumps(msg, indent=2))
# tell btest we're done
with open('done', 'w') as f:
f.write('done')
await ws.close()
sys.exit()
except:
if not connected:
print(f'failed to connect to {ws_url}, try again', file=sys.stderr)
time.sleep(1)
else:
sys.exit()
loop = asyncio.get_event_loop()
loop.run_until_complete(do_run())
@TEST-END-FILE
@TEST-START-FILE send.py
import asyncio, websockets, os, json, sys
ws_port = os.environ['BROKER_WEB_SOCKET_PORT'].split('/')[0]
ws_url = f'ws://localhost:{ws_port}/v1/messages/json'
msg_data = []
for i in range(200):
msg_data.append({
'@data-type': 'count',
'data': i
})
# Message with timestamps using the local time zone.
msg = {
'type': 'data-message',
'topic': '/test',
'@data-type': "vector",
'data': msg_data
}
async def do_run():
async with websockets.connect(ws_url) as ws:
await ws.send('[]')
await ws.recv() # wait for ACK
msg_str = json.dumps(msg, indent=2)
if len(msg_str) < 4096:
raise RuntimeError('message is too short')
# split the message into chunks of 4096 bytes
chunks = [msg_str[i:i+4096] for i in range(0, len(msg_str), 4096)]
await ws.send(chunks) # send as fragmented message
resp = await ws.recv()
print(resp)
await ws.close()
loop = asyncio.get_event_loop()
loop.run_until_complete(do_run())
@TEST-END-FILE