Skip to content

Commit 5f49275

Browse files
authored
feat: upgrade websockets to v14 to adapt to library API changes (#425)
### Description - Adopt [websockets](https://pypi.org/project/websockets/) v14. ### Issues - Closes: #325
1 parent a546e75 commit 5f49275

File tree

5 files changed

+88
-90
lines changed

5 files changed

+88
-90
lines changed

pyproject.toml

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,7 @@ dependencies = [
4242
"lazy-object-proxy>=1.10.0",
4343
"more_itertools>=10.2.0",
4444
"typing-extensions>=4.1.0",
45-
# TODO: Relax the upper bound once the issue is resolved:
46-
# https://github.com/apify/apify-sdk-python/issues/325
47-
"websockets>=10.0,<14.0.0",
45+
"websockets>=14.0",
4846
]
4947

5048
[project.optional-dependencies]

src/apify/_platform_event_manager.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from datetime import datetime
55
from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
66

7-
import websockets.client
7+
import websockets.asyncio.client
88
from pydantic import BaseModel, Discriminator, Field, TypeAdapter
99
from typing_extensions import Self, Unpack, override
1010

@@ -143,7 +143,7 @@ class PlatformEventManager(EventManager):
143143
but instead use it via the `Actor.on()` and `Actor.off()` methods.
144144
"""
145145

146-
_platform_events_websocket: websockets.client.WebSocketClientProtocol | None = None
146+
_platform_events_websocket: websockets.asyncio.client.ClientConnection | None = None
147147
_process_platform_messages_task: asyncio.Task | None = None
148148
_send_system_info_interval_task: asyncio.Task | None = None
149149
_connected_to_platform_websocket: asyncio.Future = asyncio.Future()
@@ -196,7 +196,7 @@ async def __aexit__(
196196

197197
async def _process_platform_messages(self, ws_url: str) -> None:
198198
try:
199-
async with websockets.client.connect(ws_url) as websocket:
199+
async with websockets.asyncio.client.connect(ws_url) as websocket:
200200
self._platform_events_websocket = websocket
201201
self._connected_to_platform_websocket.set_result(True)
202202

tests/unit/actor/test_actor_lifecycle.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from unittest.mock import AsyncMock, Mock
1010

1111
import pytest
12-
import websockets.server
12+
import websockets.asyncio.server
1313
from lazy_object_proxy import Proxy
1414

1515
from apify_shared.consts import ActorEnvVars, ApifyEnvVars
@@ -139,10 +139,10 @@ def log_persist_state(data: Any) -> None:
139139
nonlocal persist_state_events_data
140140
persist_state_events_data.append(data)
141141

142-
async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
142+
async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
143143
await websocket.wait_closed()
144144

145-
async with websockets.server.serve(handler, host='localhost') as ws_server:
145+
async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
146146
port: int = ws_server.sockets[0].getsockname()[1] # type: ignore[index]
147147
monkeypatch.setenv(ApifyEnvVars.ACTOR_EVENTS_WS_URL, f'ws://localhost:{port}')
148148

@@ -181,7 +181,7 @@ async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
181181
Actor.on(Event.PERSIST_STATE, log_persist_state)
182182
await asyncio.sleep(2)
183183

184-
for socket in ws_server.websockets:
184+
for socket in ws_server.connections:
185185
await socket.send(
186186
json.dumps(
187187
{

tests/unit/test_platform_event_manager.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
import pytest
1111
import websockets
12-
import websockets.server
12+
import websockets.asyncio.server
1313

1414
from apify_shared.consts import ActorEnvVars
1515
from crawlee.events._types import Event
@@ -133,16 +133,16 @@ async def test_lifecycle_on_platform_without_websocket(monkeypatch: pytest.Monke
133133

134134

135135
async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
136-
connected_ws_clients: set[websockets.server.WebSocketServerProtocol] = set()
136+
connected_ws_clients: set[websockets.asyncio.server.ServerConnection] = set()
137137

138-
async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
138+
async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
139139
connected_ws_clients.add(websocket)
140140
try:
141141
await websocket.wait_closed()
142142
finally:
143143
connected_ws_clients.remove(websocket)
144144

145-
async with websockets.server.serve(handler, host='localhost') as ws_server:
145+
async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
146146
# When you don't specify a port explicitly, the websocket connection is opened on a random free port.
147147
# We need to find out which port is that.
148148
port: int = ws_server.sockets[0].getsockname()[1] # type: ignore[index]
@@ -153,9 +153,9 @@ async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
153153

154154

155155
async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
156-
connected_ws_clients: set[websockets.server.WebSocketServerProtocol] = set()
156+
connected_ws_clients: set[websockets.asyncio.server.ServerConnection] = set()
157157

158-
async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
158+
async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
159159
connected_ws_clients.add(websocket)
160160
try:
161161
await websocket.wait_closed()
@@ -169,7 +169,7 @@ async def send_platform_event(event_name: Event, data: Any = None) -> None:
169169

170170
websockets.broadcast(connected_ws_clients, json.dumps(message))
171171

172-
async with websockets.server.serve(handler, host='localhost') as ws_server:
172+
async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
173173
# When you don't specify a port explicitly, the websocket connection is opened on a random free port.
174174
# We need to find out which port is that.
175175
port: int = ws_server.sockets[0].getsockname()[1] # type: ignore[index]

0 commit comments

Comments
 (0)