"""
Unit test file for rpc/external_message_consumer.py
"""
import asyncio
import logging
from datetime import UTC, datetime
from unittest.mock import MagicMock
import pytest
import websockets
from freqtrade.data.dataprovider import DataProvider
from freqtrade.rpc.external_message_consumer import ExternalMessageConsumer
from tests.conftest import log_has, log_has_re, log_has_when
_TEST_WS_TOKEN = "secret_Ws_t0ken"
_TEST_WS_HOST = "127.0.0.1"
_TEST_WS_PORT = 9989
@pytest.fixture
def patched_emc(default_conf, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{"name": "default", "host": "null", "port": 9891, "ws_token": _TEST_WS_TOKEN}
],
}
}
)
dataprovider = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dataprovider)
try:
yield emc
finally:
emc.shutdown()
def test_emc_start(patched_emc, caplog):
# Test if the message was printed
assert log_has_when("Starting ExternalMessageConsumer", caplog, "setup")
# Test if the thread and loop objects were created
assert patched_emc._thread and patched_emc._loop
# Test we call start again nothing happens
prev_thread = patched_emc._thread
patched_emc.start()
assert prev_thread == patched_emc._thread
def test_emc_shutdown(patched_emc, caplog):
patched_emc.shutdown()
assert log_has("Stopping ExternalMessageConsumer", caplog)
# Test the loop has stopped
assert patched_emc._loop is None
# Test if the thread has stopped
assert patched_emc._thread is None
caplog.clear()
patched_emc.shutdown()
# Test func didn't run again as it was called once already
assert not log_has("Stopping ExternalMessageConsumer", caplog)
def test_emc_init(patched_emc):
# Test the settings were set correctly
assert patched_emc.initial_candle_limit <= 1500
assert patched_emc.wait_timeout > 0
assert patched_emc.sleep_time > 0
# Parametrize this?
def test_emc_handle_producer_message(patched_emc, caplog, ohlcv_history):
test_producer = {"name": "test", "url": "ws://test", "ws_token": "test"}
producer_name = test_producer["name"]
invalid_msg = r"Invalid message .+"
caplog.set_level(logging.DEBUG)
# Test handle whitelist message
whitelist_message = {"type": "whitelist", "data": ["BTC/USDT"]}
patched_emc.handle_producer_message(test_producer, whitelist_message)
assert log_has(f"Received message of type `whitelist` from `{producer_name}`", caplog)
assert log_has(
f"Consumed message from `{producer_name}` of type `RPCMessageType.WHITELIST`", caplog
)
# Test handle analyzed_df single candle message
df_message = {
"type": "analyzed_df",
"data": {
"key": ("BTC/USDT", "5m", "spot"),
"df": ohlcv_history,
"la": datetime.now(UTC),
},
}
patched_emc.handle_producer_message(test_producer, df_message)
assert log_has(f"Received message of type `analyzed_df` from `{producer_name}`", caplog)
assert log_has_re(r"Holes in data or no existing df,.+", caplog)
# Test unhandled message
unhandled_message = {"type": "status", "data": "RUNNING"}
patched_emc.handle_producer_message(test_producer, unhandled_message)
assert log_has_re(r"Received unhandled message\: .*", caplog)
# Test malformed messages
caplog.clear()
malformed_message = {"type": "whitelist", "data": {"pair": "BTC/USDT"}}
patched_emc.handle_producer_message(test_producer, malformed_message)
assert log_has_re(invalid_msg, caplog)
caplog.clear()
malformed_message = {
"type": "analyzed_df",
"data": {"key": "BTC/USDT", "df": ohlcv_history, "la": datetime.now(UTC)},
}
patched_emc.handle_producer_message(test_producer, malformed_message)
assert log_has(f"Received message of type `analyzed_df` from `{producer_name}`", caplog)
assert log_has_re(invalid_msg, caplog)
caplog.clear()
# Empty dataframe
malformed_message = {
"type": "analyzed_df",
"data": {
"key": ("BTC/USDT", "5m", "spot"),
"df": ohlcv_history.loc[ohlcv_history["open"] < 0],
"la": datetime.now(UTC),
},
}
patched_emc.handle_producer_message(test_producer, malformed_message)
assert log_has(f"Received message of type `analyzed_df` from `{producer_name}`", caplog)
assert not log_has_re(invalid_msg, caplog)
assert log_has_re(r"Received Empty Dataframe for.+", caplog)
caplog.clear()
malformed_message = {"some": "stuff"}
patched_emc.handle_producer_message(test_producer, malformed_message)
assert log_has_re(invalid_msg, caplog)
caplog.clear()
caplog.clear()
malformed_message = {"type": "whitelist", "data": None}
patched_emc.handle_producer_message(test_producer, malformed_message)
assert log_has_re(r"Empty message .+", caplog)
async def test_emc_create_connection_success(default_conf, caplog, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 60,
"ping_timeout": 60,
"sleep_timeout": 60,
}
}
)
mocker.patch(
"freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start", MagicMock()
)
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
test_producer = default_conf["external_message_consumer"]["producers"][0]
lock = asyncio.Lock()
emc._running = True
async def eat(websocket):
emc._running = False
try:
async with websockets.serve(eat, _TEST_WS_HOST, _TEST_WS_PORT):
await emc._create_connection(test_producer, lock)
assert log_has_re(r"Connected to channel.+", caplog)
finally:
emc.shutdown()
@pytest.mark.parametrize(
"host,port",
[
(_TEST_WS_HOST, -1),
("10000.1241..2121/", _TEST_WS_PORT),
],
)
async def test_emc_create_connection_invalid_url(default_conf, caplog, mocker, host, port):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{"name": "default", "host": host, "port": port, "ws_token": _TEST_WS_TOKEN}
],
"wait_timeout": 60,
"ping_timeout": 60,
"sleep_timeout": 60,
}
}
)
dp = DataProvider(default_conf, None, None, None)
# Handle start explicitly to avoid messing with threading in tests
mocker.patch("freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start")
mocker.patch("freqtrade.rpc.api_server.ws.channel.create_channel")
emc = ExternalMessageConsumer(default_conf, dp)
try:
emc._running = True
await emc._create_connection(emc.producers[0], asyncio.Lock())
assert log_has_re(r".+ is an invalid WebSocket URL .+", caplog)
finally:
emc.shutdown()
async def test_emc_create_connection_error(default_conf, caplog, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 60,
"ping_timeout": 60,
"sleep_timeout": 60,
}
}
)
# Test unexpected error
mocker.patch("websockets.connect", side_effect=RuntimeError)
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
try:
await asyncio.sleep(0.05)
assert log_has("Unexpected error has occurred:", caplog)
finally:
emc.shutdown()
async def test_emc_receive_messages_valid(default_conf, caplog, mocker):
caplog.set_level(logging.DEBUG)
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 1,
"ping_timeout": 60,
"sleep_time": 60,
}
}
)
mocker.patch(
"freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start", MagicMock()
)
lock = asyncio.Lock()
test_producer = default_conf["external_message_consumer"]["producers"][0]
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
class TestChannel:
async def recv(self, *args, **kwargs):
emc._running = False
return {"type": "whitelist", "data": ["BTC/USDT"]}
async def ping(self, *args, **kwargs):
return asyncio.Future()
try:
emc._running = True
await emc._receive_messages(TestChannel(), test_producer, lock)
assert log_has_re(r"Received message of type `whitelist`.+", caplog)
finally:
emc.shutdown()
async def test_emc_receive_messages_invalid(default_conf, caplog, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 1,
"ping_timeout": 60,
"sleep_time": 60,
}
}
)
mocker.patch(
"freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start", MagicMock()
)
lock = asyncio.Lock()
test_producer = default_conf["external_message_consumer"]["producers"][0]
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
class TestChannel:
async def recv(self, *args, **kwargs):
emc._running = False
return {"type": ["BTC/USDT"]}
async def ping(self, *args, **kwargs):
return asyncio.Future()
try:
emc._running = True
await emc._receive_messages(TestChannel(), test_producer, lock)
assert log_has_re(r"Invalid message from.+", caplog)
finally:
emc.shutdown()
async def test_emc_receive_messages_timeout(default_conf, caplog, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 0.1,
"ping_timeout": 1,
"sleep_time": 1,
}
}
)
mocker.patch(
"freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start", MagicMock()
)
lock = asyncio.Lock()
test_producer = default_conf["external_message_consumer"]["producers"][0]
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
def change_running():
emc._running = not emc._running
class TestChannel:
async def recv(self, *args, **kwargs):
await asyncio.sleep(0.2)
async def ping(self, *args, **kwargs):
return asyncio.Future()
try:
change_running()
with pytest.raises(asyncio.TimeoutError):
await emc._receive_messages(TestChannel(), test_producer, lock)
assert log_has_re(r"Ping error.+", caplog)
finally:
emc.shutdown()
async def test_emc_receive_messages_handle_error(default_conf, caplog, mocker):
default_conf.update(
{
"external_message_consumer": {
"enabled": True,
"producers": [
{
"name": "default",
"host": _TEST_WS_HOST,
"port": _TEST_WS_PORT,
"ws_token": _TEST_WS_TOKEN,
}
],
"wait_timeout": 1,
"ping_timeout": 1,
"sleep_time": 1,
}
}
)
mocker.patch(
"freqtrade.rpc.external_message_consumer.ExternalMessageConsumer.start", MagicMock()
)
lock = asyncio.Lock()
test_producer = default_conf["external_message_consumer"]["producers"][0]
dp = DataProvider(default_conf, None, None, None)
emc = ExternalMessageConsumer(default_conf, dp)
emc.handle_producer_message = MagicMock(side_effect=Exception)
class TestChannel:
async def recv(self, *args, **kwargs):
emc._running = False
return {"type": "whitelist", "data": ["BTC/USDT"]}
async def ping(self, *args, **kwargs):
return asyncio.Future()
try:
emc._running = True
await emc._receive_messages(TestChannel(), test_producer, lock)
assert log_has_re(r"Error handling producer message.+", caplog)
finally:
emc.shutdown()