From cead71c1fe89da6b17875938df8f700d6a8d4e30 Mon Sep 17 00:00:00 2001
From: Quitterie Lucas <quitterielucas@outlook.fr>
Date: Fri, 20 Sep 2024 17:13:35 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7(backends)=20fix=20tests=20failure?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

WIP.
---
 src/ralph/backends/data/async_ws.py   |  7 +++----
 tests/backends/data/test_async_lrs.py |  1 +
 tests/backends/test_loader.py         | 22 ++++++++++++----------
 tests/fixtures/backends.py            |  4 +++-
 tests/test_cli.py                     |  2 +-
 5 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/src/ralph/backends/data/async_ws.py b/src/ralph/backends/data/async_ws.py
index 55dbbf7f4..6f2c495c3 100644
--- a/src/ralph/backends/data/async_ws.py
+++ b/src/ralph/backends/data/async_ws.py
@@ -7,6 +7,7 @@
 from pydantic import AnyUrl, PositiveInt
 from pydantic_settings import SettingsConfigDict
 from websockets.http11 import USER_AGENT
+from websockets.asyncio.client import connect, ClientConnection
 
 from ralph.backends.data.base import (
     BaseAsyncDataBackend,
@@ -40,7 +41,6 @@ class WSClientOptions(ClientOptions):
             Setting it to `None` disables keepalive pings.
         ping_timeout (float): Timeout for keepalive pings in seconds.
             Setting it to `None` disables timeouts.
-        read_limit (int): High-water mark of read buffer in bytes.
         user_agent_header (str): Value of  the `User-Agent` request header.
             It defaults to "Python/x.y.z websockets/X.Y".
             Setting it to `None` removes the header.
@@ -55,7 +55,6 @@ class WSClientOptions(ClientOptions):
     origin: Optional[str] = None
     ping_interval: Optional[float] = 20
     ping_timeout: Optional[float] = 20
-    read_limit: int = 2**16
     user_agent_header: Optional[str] = USER_AGENT
     write_limit: int = 2**16
 
@@ -93,11 +92,11 @@ def __init__(self, settings: Optional[WSDataBackendSettings] = None):
         super().__init__(settings)
         self._client = None
 
-    async def client(self) -> websockets.WebSocketClientProtocol:
+    async def client(self) -> ClientConnection:
         """Create a websocket client connected to `settings.URI` if it doesn't exist."""
         if not self._client:
             try:
-                self._client = await websockets.connect(
+                self._client = await connect(
                     str(self.settings.URI), **self.settings.CLIENT_OPTIONS.model_dump()
                 )
             except (websockets.WebSocketException, OSError, TimeoutError) as error:
diff --git a/tests/backends/data/test_async_lrs.py b/tests/backends/data/test_async_lrs.py
index 5ac9e7894..6e91a87e0 100644
--- a/tests/backends/data/test_async_lrs.py
+++ b/tests/backends/data/test_async_lrs.py
@@ -809,3 +809,4 @@ async def simulate_network_latency(_):
     # Server side processing time should be 3 times faster with unlimited
     assert limited_concurrency_duration > 2.1 * concurrent_duration
     assert limited_concurrency_duration <= 3.1 * concurrent_duration
+
diff --git a/tests/backends/test_loader.py b/tests/backends/test_loader.py
index 9c0816b09..f9644ce45 100644
--- a/tests/backends/test_loader.py
+++ b/tests/backends/test_loader.py
@@ -189,14 +189,16 @@ def mock_entry_points(group):
             ),
         ]
 
-    monkeypatch.setattr("ralph.backends.loader.entry_points", mock_entry_points)
+    with monkeypatch.context() as mocked_context:
+        mocked_context.setattr("ralph.backends.loader.entry_points", mock_entry_points)
+        get_lrs_backends.cache_clear()
+        assert get_lrs_backends() == {
+            "test_backend": TestBackend,
+            "async_es": AsyncESLRSBackend,
+            "async_mongo": AsyncMongoLRSBackend,
+            "clickhouse": ClickHouseLRSBackend,
+            "es": ESLRSBackend,
+            "fs": FSLRSBackend,
+            "mongo": MongoLRSBackend,
+        }
     get_lrs_backends.cache_clear()
-    assert get_lrs_backends() == {
-        "test_backend": TestBackend,
-        "async_es": AsyncESLRSBackend,
-        "async_mongo": AsyncMongoLRSBackend,
-        "clickhouse": ClickHouseLRSBackend,
-        "es": ESLRSBackend,
-        "fs": FSLRSBackend,
-        "mongo": MongoLRSBackend,
-    }
diff --git a/tests/fixtures/backends.py b/tests/fixtures/backends.py
index 8c95ca997..d03ff22a4 100644
--- a/tests/fixtures/backends.py
+++ b/tests/fixtures/backends.py
@@ -22,6 +22,8 @@
 from pymongo import MongoClient
 from pymongo.errors import CollectionInvalid
 
+from websockets.asyncio.server import serve
+
 from ralph.backends.data.async_es import AsyncESDataBackend
 from ralph.backends.data.async_lrs import AsyncLRSDataBackend
 from ralph.backends.data.async_mongo import AsyncMongoDataBackend
@@ -893,7 +895,7 @@ async def forward(websocket):
             await websocket.send(json.dumps(event))
             await asyncio.sleep(random.randrange(0, 500) / 10000.0)
 
-    async with websockets.serve(forward, "0.0.0.0", WS_TEST_PORT) as server:
+    async with serve(forward, "0.0.0.0", WS_TEST_PORT) as server:
         yield server
 
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index a3ebd85b2..b6dfac269 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -6,7 +6,7 @@
 from importlib import reload
 from pathlib import Path
 from typing import Optional, Union
-
+import asyncio
 import pytest
 from click.exceptions import BadParameter
 from click.testing import CliRunner