freqtrade_origin/tests/rpc/test_rpc_emc.py

"""
Unit test file for rpc/external_message_consumer.py
"""

import asyncio
import logging
from datetime import datetime, timezone
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(timezone.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(timezone.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(timezone.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()