Merge pull request #7545 from wizrds/feat/test-ws-client
Message WebSocket Testing client
This commit is contained in:
Matthias
GitHub
commit
1cabfe8d0a
321
scripts/ws_client.py
Normal file
321
scripts/ws_client.py
Normal file
@ -0,0 +1,321 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Simple command line client for Testing/debugging
|
||||
a Freqtrade bot's message websocket
|
||||
|
||||
Should not import anything from freqtrade,
|
||||
so it can be used as a standalone script.
|
||||
"""
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import orjson
|
||||
import pandas
|
||||
import rapidjson
|
||||
import websockets
|
||||
from dateutil.relativedelta import relativedelta
|
||||
|
||||
|
||||
logger = logging.getLogger("WebSocketClient")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def setup_logging(filename: str):
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
handlers=[
|
||||
logging.FileHandler(filename),
|
||||
logging.StreamHandler()
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'-c',
|
||||
'--config',
|
||||
help='Specify configuration file (default: %(default)s). ',
|
||||
dest='config',
|
||||
type=str,
|
||||
metavar='PATH',
|
||||
default='config.json'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-l',
|
||||
'--logfile',
|
||||
help='The filename to log to.',
|
||||
dest='logfile',
|
||||
type=str,
|
||||
default='ws_client.log'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return vars(args)
|
||||
|
||||
|
||||
def load_config(configfile):
|
||||
file = Path(configfile)
|
||||
if file.is_file():
|
||||
with file.open("r") as f:
|
||||
config = rapidjson.load(f, parse_mode=rapidjson.PM_COMMENTS |
|
||||
rapidjson.PM_TRAILING_COMMAS)
|
||||
return config
|
||||
else:
|
||||
logger.warning(f"Could not load config file {file}.")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def readable_timedelta(delta):
|
||||
"""
|
||||
Convert a dateutil.relativedelta to a readable format
|
||||
|
||||
:param delta: A dateutil.relativedelta
|
||||
:returns: The readable time difference string
|
||||
"""
|
||||
attrs = ['years', 'months', 'days', 'hours', 'minutes', 'seconds', 'microseconds']
|
||||
return ", ".join([
|
||||
'%d %s' % (getattr(delta, attr), attr if getattr(delta, attr) > 1 else attr[:-1])
|
||||
for attr in attrs if getattr(delta, attr)
|
||||
])
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def json_serialize(message):
|
||||
"""
|
||||
Serialize a message to JSON using orjson
|
||||
:param message: The message to serialize
|
||||
"""
|
||||
return str(orjson.dumps(message), "utf-8")
|
||||
|
||||
|
||||
def json_deserialize(message):
|
||||
"""
|
||||
Deserialize JSON to a dict
|
||||
:param message: The message to deserialize
|
||||
"""
|
||||
def json_to_dataframe(data: str) -> pandas.DataFrame:
|
||||
dataframe = pandas.read_json(data, orient='split')
|
||||
if 'date' in dataframe.columns:
|
||||
dataframe['date'] = pandas.to_datetime(dataframe['date'], unit='ms', utc=True)
|
||||
|
||||
return dataframe
|
||||
|
||||
def _json_object_hook(z):
|
||||
if z.get('__type__') == 'dataframe':
|
||||
return json_to_dataframe(z.get('__value__'))
|
||||
return z
|
||||
|
||||
return rapidjson.loads(message, object_hook=_json_object_hook)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ClientProtocol:
|
||||
logger = logging.getLogger("WebSocketClient.Protocol")
|
||||
_MESSAGE_COUNT = 0
|
||||
_LAST_RECEIVED_AT = 0 # The epoch we received a message most recently
|
||||
|
||||
async def on_connect(self, websocket):
|
||||
# On connection we have to send our initial requests
|
||||
initial_requests = [
|
||||
{
|
||||
"type": "subscribe", # The subscribe request should always be first
|
||||
"data": ["analyzed_df", "whitelist"] # The message types we want
|
||||
},
|
||||
{
|
||||
"type": "whitelist",
|
||||
"data": None,
|
||||
},
|
||||
{
|
||||
"type": "analyzed_df",
|
||||
"data": {"limit": 1500}
|
||||
}
|
||||
]
|
||||
|
||||
for request in initial_requests:
|
||||
await websocket.send(json_serialize(request))
|
||||
|
||||
async def on_message(self, websocket, name, message):
|
||||
deserialized = json_deserialize(message)
|
||||
|
||||
message_size = sys.getsizeof(message)
|
||||
message_type = deserialized.get('type')
|
||||
message_data = deserialized.get('data')
|
||||
|
||||
self.logger.info(
|
||||
f"Received message of type {message_type} [{message_size} bytes] @ [{name}]"
|
||||
)
|
||||
|
||||
time_difference = self._calculate_time_difference()
|
||||
|
||||
if self._MESSAGE_COUNT > 0:
|
||||
self.logger.info(f"Time since last message: {time_difference}")
|
||||
|
||||
message_handler = getattr(self, f"_handle_{message_type}", None) or self._handle_default
|
||||
await message_handler(name, message_type, message_data)
|
||||
|
||||
self._MESSAGE_COUNT += 1
|
||||
self.logger.info(f"[{self._MESSAGE_COUNT}] total messages..")
|
||||
self.logger.info("-" * 80)
|
||||
|
||||
def _calculate_time_difference(self):
|
||||
old_last_received_at = self._LAST_RECEIVED_AT
|
||||
self._LAST_RECEIVED_AT = time.time()
|
||||
time_delta = relativedelta(seconds=(self._LAST_RECEIVED_AT - old_last_received_at))
|
||||
|
||||
return readable_timedelta(time_delta)
|
||||
|
||||
async def _handle_whitelist(self, name, type, data):
|
||||
self.logger.info(data)
|
||||
|
||||
async def _handle_default(self, name, type, data):
|
||||
key, la, df = data['key'], data['la'], data['df']
|
||||
|
||||
if not df.empty:
|
||||
columns = ", ".join([str(column) for column in df.columns])
|
||||
|
||||
self.logger.info(key)
|
||||
self.logger.info(f"Last analyzed datetime: {la}")
|
||||
self.logger.info(f"Latest candle datetime: {df.iloc[-1]['date']}")
|
||||
self.logger.info(f"DataFrame length: {len(df)}")
|
||||
self.logger.info(f"DataFrame columns: {columns}")
|
||||
|
||||
|
||||
async def create_client(
|
||||
host,
|
||||
port,
|
||||
token,
|
||||
name='default',
|
||||
protocol=ClientProtocol(),
|
||||
sleep_time=10,
|
||||
ping_timeout=10,
|
||||
wait_timeout=30,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Create a websocket client and listen for messages
|
||||
:param host: The host
|
||||
:param port: The port
|
||||
:param token: The websocket auth token
|
||||
:param name: The name of the producer
|
||||
:param **kwargs: Any extra kwargs passed to websockets.connect
|
||||
"""
|
||||
|
||||
while 1:
|
||||
try:
|
||||
websocket_url = f"ws://{host}:{port}/api/v1/message/ws?token={token}"
|
||||
logger.info(f"Attempting to connect to {name} @ {host}:{port}")
|
||||
|
||||
async with websockets.connect(websocket_url, **kwargs) as ws:
|
||||
logger.info("Connection successful...")
|
||||
await protocol.on_connect(ws)
|
||||
|
||||
# Now listen for messages
|
||||
while 1:
|
||||
try:
|
||||
message = await asyncio.wait_for(
|
||||
ws.recv(),
|
||||
timeout=wait_timeout
|
||||
)
|
||||
|
||||
await protocol.on_message(ws, name, message)
|
||||
|
||||
except (
|
||||
asyncio.TimeoutError,
|
||||
websockets.exceptions.ConnectionClosed
|
||||
):
|
||||
# Try pinging
|
||||
try:
|
||||
pong = ws.ping()
|
||||
await asyncio.wait_for(
|
||||
pong,
|
||||
timeout=ping_timeout
|
||||
)
|
||||
logger.info("Connection still alive...")
|
||||
|
||||
continue
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(f"Ping timed out, retrying in {sleep_time}s")
|
||||
await asyncio.sleep(sleep_time)
|
||||
|
||||
break
|
||||
|
||||
except (
|
||||
socket.gaierror,
|
||||
ConnectionRefusedError,
|
||||
websockets.exceptions.InvalidStatusCode,
|
||||
websockets.exceptions.InvalidMessage
|
||||
) as e:
|
||||
logger.error(f"Connection Refused - {e} retrying in {sleep_time}s")
|
||||
await asyncio.sleep(sleep_time)
|
||||
|
||||
continue
|
||||
|
||||
except (
|
||||
websockets.exceptions.ConnectionClosedError,
|
||||
websockets.exceptions.ConnectionClosedOK
|
||||
):
|
||||
# Just keep trying to connect again indefinitely
|
||||
await asyncio.sleep(sleep_time)
|
||||
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
# An unforseen error has occurred, log and try reconnecting again
|
||||
logger.error("Unexpected error has occurred:")
|
||||
logger.exception(e)
|
||||
|
||||
await asyncio.sleep(sleep_time)
|
||||
continue
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _main(args):
|
||||
setup_logging(args['logfile'])
|
||||
config = load_config(args['config'])
|
||||
|
||||
emc_config = config.get('external_message_consumer', {})
|
||||
|
||||
producers = emc_config.get('producers', [])
|
||||
producer = producers[0]
|
||||
|
||||
wait_timeout = emc_config.get('wait_timeout', 300)
|
||||
ping_timeout = emc_config.get('ping_timeout', 10)
|
||||
sleep_time = emc_config.get('sleep_time', 10)
|
||||
message_size_limit = (emc_config.get('message_size_limit', 8) << 20)
|
||||
|
||||
await create_client(
|
||||
producer['host'],
|
||||
producer['port'],
|
||||
producer['ws_token'],
|
||||
producer['name'],
|
||||
sleep_time=sleep_time,
|
||||
ping_timeout=ping_timeout,
|
||||
wait_timeout=wait_timeout,
|
||||
max_size=message_size_limit
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
try:
|
||||
asyncio.run(_main(args))
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Exiting...")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user