moved ws_schemas, first ws tests

freqtrade/rpc/api_server/api_ws.py

@@ -2,6 +2,7 @@ import logging
 from typing import Any, Dict
 
 from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
+from pydantic import ValidationError
 # fastapi does not make this available through it, so import directly from starlette
 from starlette.websockets import WebSocketState
 
@@ -9,9 +10,8 @@ from freqtrade.enums import RPCMessageType, RPCRequestType
 from freqtrade.rpc.api_server.api_auth import get_ws_token
 from freqtrade.rpc.api_server.deps import get_channel_manager, get_rpc
 from freqtrade.rpc.api_server.ws.channel import WebSocketChannel
-from freqtrade.rpc.api_server.ws.schema import (ValidationError, WSAnalyzedDFMessage,
-                                                WSMessageSchema, WSRequestSchema,
-                                                WSWhitelistMessage)
+from freqtrade.rpc.api_server.ws_schemas import (WSAnalyzedDFMessage, WSMessageSchema,
+                                                 WSRequestSchema, WSWhitelistMessage)
 from freqtrade.rpc.rpc import RPC
 
 
@@ -102,13 +102,11 @@ async def message_endpoint(
     Message WebSocket endpoint, facilitates sending RPC messages
     """
     try:
-        if is_websocket_alive(ws):
-            # TODO:
-            # Return a channel ID, pass that instead of ws to the rest of the methods
-            channel = await channel_manager.on_connect(ws)
+        # TODO:
+        # Return a channel ID, pass that instead of ws to the rest of the methods
+        channel = await channel_manager.on_connect(ws)
 
-            if not channel:
-                return
+        if await is_websocket_alive(ws):
 
             logger.info(f"Consumer connected - {channel}")
 
@@ -131,6 +129,9 @@ async def message_endpoint(
                 # RuntimeError('Cannot call "send" once a closed message has been sent')
                 await channel_manager.on_disconnect(ws)
 
+        else:
+            ws.close()
+
     except Exception as e:
         logger.error(f"Failed to serve - {ws.client}")
         # Log tracebacks to keep track of what errors are happening

freqtrade/rpc/api_server/ws_schemas.py

@@ -2,15 +2,12 @@ from datetime import datetime
 from typing import Any, Dict, List, Optional
 
 from pandas import DataFrame
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel
 
 from freqtrade.constants import PairWithTimeframe
 from freqtrade.enums.rpcmessagetype import RPCMessageType, RPCRequestType
 
 
-__all__ = ('WSRequestSchema', 'WSMessageSchema', 'ValidationError')
-
-
 class BaseArbitraryModel(BaseModel):
     class Config:
         arbitrary_types_allowed = True

freqtrade/rpc/external_message_consumer.py

@@ -11,15 +11,16 @@ from threading import Thread
 from typing import TYPE_CHECKING, Any, Dict, List
 
 import websockets
+from pydantic import ValidationError
 
 from freqtrade.data.dataprovider import DataProvider
 from freqtrade.enums import RPCMessageType
 from freqtrade.misc import remove_entry_exit_signals
 from freqtrade.rpc.api_server.ws.channel import WebSocketChannel
-from freqtrade.rpc.api_server.ws.schema import (ValidationError, WSAnalyzedDFMessage,
-                                                WSAnalyzedDFRequest, WSMessageSchema,
-                                                WSRequestSchema, WSSubscribeRequest,
-                                                WSWhitelistMessage, WSWhitelistRequest)
+from freqtrade.rpc.api_server.ws_schemas import (WSAnalyzedDFMessage, WSAnalyzedDFRequest,
+                                                 WSMessageSchema, WSRequestSchema,
+                                                 WSSubscribeRequest, WSWhitelistMessage,
+                                                 WSWhitelistRequest)
 
 
 if TYPE_CHECKING: