apify 2.4.0b2__tar.gz → 2.4.0b4__tar.gz

{apify-2.4.0b2 → apify-2.4.0b4}/CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 ### 🚀 Features
 
 - Update to Crawlee v0.6 ([#420](https://github.com/apify/apify-sdk-python/pull/420)) ([9be4336](https://github.com/apify/apify-sdk-python/commit/9be433667231cc5739861fa693d7a726860d6aca)) by [@vdusek](https://github.com/vdusek)
+- Add Actor `exit_process` option ([#424](https://github.com/apify/apify-sdk-python/pull/424)) ([994c832](https://github.com/apify/apify-sdk-python/commit/994c8323b994e009db0ccdcb624891a2fef97070)) by [@vdusek](https://github.com/vdusek), closes [#396](https://github.com/apify/apify-sdk-python/issues/396), [#401](https://github.com/apify/apify-sdk-python/issues/401)
 
 
 <!-- git-cliff-unreleased-end -->

{apify-2.4.0b2 → apify-2.4.0b4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apify
-Version: 2.4.0b2
+Version: 2.4.0b4
 Summary: Apify SDK for Python
 Project-URL: Homepage, https://docs.apify.com/sdk/python/
 Project-URL: Apify homepage, https://apify.com
@@ -231,7 +231,7 @@ Requires-Dist: httpx>=0.27.0
 Requires-Dist: lazy-object-proxy>=1.10.0
 Requires-Dist: more-itertools>=10.2.0
 Requires-Dist: typing-extensions>=4.1.0
-Requires-Dist: websockets<14.0.0,>=10.0
+Requires-Dist: websockets>=14.0
 Provides-Extra: scrapy
 Requires-Dist: scrapy>=2.11.0; extra == 'scrapy'
 Description-Content-Type: text/markdown

apify-2.4.0b4/docs/03_concepts/12_pay_per_event.mdx

@@ -0,0 +1,46 @@
+---
+id: pay-per-event
+title: Pay-per-event monetization
+description: Monetize your Actors using the pay-per-event pricing model
+---
+
+import ActorChargeSource from '!!raw-loader!./code/actor_charge.py';
+import ConditionalActorChargeSource from '!!raw-loader!./code/conditional_actor_charge.py';
+import ApiLink from '@site/src/components/ApiLink';
+import CodeBlock from '@theme/CodeBlock';
+
+Apify provides several [pricing models](https://docs.apify.com/platform/actors/publishing/monetize) for monetizing your Actors. The most recent and most flexible one is [pay-per-event](https://docs.apify.com/platform/actors/running/actors-in-store#pay-per-event), which lets you charge your users programmatically directly from your Actor. As the name suggests, you may charge the users each time a specific event occurs, for example a call to an external API or when you return a result.
+
+To use the pay-per-event pricing model, you first need to [set it up](https://docs.apify.com/platform/actors/running/actors-in-store#pay-per-event) for your Actor in the Apify console. After that, you're free to start charging for events.
+
+## Charging for events
+
+After monetization is set in the Apify console, you can add <ApiLink to="class/Actor#charge">`Actor.charge`</ApiLink> calls to your code and start monetizing!
+
+<CodeBlock language="python">
+{ActorChargeSource}
+</CodeBlock>
+
+Then you just push your code to Apify and that's it! The SDK will even keep track of the max total charge setting for you, so you will not provide more value than what the user chose to pay for.
+
+If you need finer control over charging, you can access call <ApiLink to="class/Actor#get_charging_manager">`Actor.get_charging_manager()`</ApiLink> to access the <ApiLink to="class/ChargingManager">`ChargingManager`</ApiLink>, which can provide more detailed information - for example how many events of each type can be charged before reaching the configured limit.
+
+## Transitioning from a different pricing model
+
+When you plan to start using the pay-per-event pricing model for an Actor that is already monetized with a different pricing model, your source code will need support both pricing models during the transition period enforced by the Apify platform. Arguably the most frequent case is the transition from the pay-per-result model which utilizes the `ACTOR_MAX_PAID_DATASET_ITEMS` environment variable to prevent returning unpaid dataset items. The following is an example how to handle such scenarios. The key part is the <ApiLink to="class/ChargingManager#get_pricing_info">`ChargingManager.get_pricing_info()`</ApiLink> method which returns information about the current pricing model.
+
+<CodeBlock language="python">
+{ConditionalActorChargeSource}
+</CodeBlock>
+
+## Local development
+
+It is encouraged to test your monetization code on your machine before releasing it to the public. To tell your Actor that it should work in pay-per-event mode, pass it the `ACTOR_TEST_PAY_PER_EVENT` environment variable:
+
+```shell
+ACTOR_TEST_PAY_PER_EVENT=true python -m youractor
+```
+
+If you also wish to see a log of all the events charged throughout the run, the Apify SDK keeps a log of charged events in a so called charging dataset. Your charging dataset can be found under the `charging_log` name (unless you change your storage settings, this dataset is stored in `storage/datasets/charging_log/`). Please note that this log is not available when running the Actor in production on the Apify platform.
+
+Because pricing configuration is stored by the Apify platform, all events will have a default price of $1.

apify-2.4.0b4/docs/03_concepts/code/actor_charge.py

@@ -0,0 +1,30 @@
+from apify import Actor
+
+
+async def main() -> None:
+    async with Actor:
+        # highlight-start
+        # Charge for a single occurence of an event
+        await Actor.charge(event_name='init')
+        # highlight-end
+
+        # Prepare some mock results
+        result = [
+            {'word': 'Lorem'},
+            {'word': 'Ipsum'},
+            {'word': 'Dolor'},
+            {'word': 'Sit'},
+            {'word': 'Amet'},
+        ]
+        # highlight-start
+        # Shortcut for charging for each pushed dataset item
+        await Actor.push_data(result, 'result-item')
+        # highlight-end
+
+        # highlight-start
+        # Or you can charge for a given number of events manually
+        await Actor.charge(
+            event_name='result-item',
+            count=len(result),
+        )
+        # highlight-end

apify-2.4.0b4/docs/03_concepts/code/conditional_actor_charge.py

@@ -0,0 +1,18 @@
+from apify import Actor
+
+
+async def main() -> None:
+    async with Actor:
+        # Check the dataset because there might already be items
+        # if the run migrated or was restarted
+        default_dataset = await Actor.open_dataset()
+        dataset_info = await default_dataset.get_info()
+        charged_items = dataset_info.item_count if dataset_info else 0
+
+        # highlight-start
+        if Actor.get_charging_manager().get_pricing_info().is_pay_per_event:
+            # highlight-end
+            await Actor.push_data({'hello': 'world'}, 'dataset-item')
+        elif charged_items < (Actor.config.max_paid_dataset_items or 0):
+            await Actor.push_data({'hello': 'world'})
+            charged_items += 1

{apify-2.4.0b2 → apify-2.4.0b4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "apify"
-version = "2.4.0b2"
+version = "2.4.0b4"
 description = "Apify SDK for Python"
 authors = [{ name = "Apify Technologies s.r.o.", email = "support@apify.com" }]
 license = { file = "LICENSE" }
@@ -42,9 +42,7 @@ dependencies = [
     "lazy-object-proxy>=1.10.0",
     "more_itertools>=10.2.0",
     "typing-extensions>=4.1.0",
-    # TODO: Relax the upper bound once the issue is resolved:
-    # https://github.com/apify/apify-sdk-python/issues/325
-    "websockets>=10.0,<14.0.0",
+    "websockets>=14.0",
 ]
 
 [project.optional-dependencies]

{apify-2.4.0b2 → apify-2.4.0b4}/src/apify/_actor.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import os
 import sys
+from contextlib import suppress
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, cast, overload
 
@@ -64,6 +65,7 @@ class _ActorType:
         configuration: Configuration | None = None,
         *,
         configure_logging: bool = True,
+        exit_process: bool | None = None,
     ) -> None:
         """Create an Actor instance.
 
@@ -74,7 +76,10 @@ class _ActorType:
             configuration: The Actor configuration to be used. If not passed, a new Configuration instance will
                 be created.
             configure_logging: Should the default logging configuration be configured?
+            exit_process: Whether the Actor should call `sys.exit` when the context manager exits. The default is
+                True except for the IPython, Pytest and Scrapy environments.
         """
+        self._exit_process = self._get_default_exit_process() if exit_process is None else exit_process
         self._is_exiting = False
 
         self._configuration = configuration or Configuration.get_global_configuration()
@@ -141,9 +146,19 @@ class _ActorType:
 
         return super().__repr__()
 
-    def __call__(self, configuration: Configuration | None = None, *, configure_logging: bool = True) -> Self:
+    def __call__(
+        self,
+        configuration: Configuration | None = None,
+        *,
+        configure_logging: bool = True,
+        exit_process: bool | None = None,
+    ) -> Self:
         """Make a new Actor instance with a non-default configuration."""
-        return self.__class__(configuration=configuration, configure_logging=configure_logging)
+        return self.__class__(
+            configuration=configuration,
+            configure_logging=configure_logging,
+            exit_process=exit_process,
+        )
 
     @property
     def apify_client(self) -> ApifyClientAsync:
@@ -281,13 +296,7 @@ class _ActorType:
         await asyncio.wait_for(finalize(), cleanup_timeout.total_seconds())
         self._is_initialized = False
 
-        if is_running_in_ipython():
-            self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running in IPython')
-        elif os.getenv('PYTEST_CURRENT_TEST', default=False):  # noqa: PLW1508
-            self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running in an unit test')
-        elif os.getenv('SCRAPY_SETTINGS_MODULE'):
-            self.log.debug(f'Not calling sys.exit({exit_code}) because Actor is running with Scrapy')
-        else:
+        if self._exit_process:
             sys.exit(exit_code)
 
     async def fail(
@@ -1128,6 +1137,26 @@ class _ActorType:
 
         return proxy_configuration
 
+    def _get_default_exit_process(self) -> bool:
+        """Returns False for IPython, Pytest, and Scrapy environments, True otherwise."""
+        if is_running_in_ipython():
+            self.log.debug('Running in IPython, setting default `exit_process` to False.')
+            return False
+
+        # Check if running in Pytest by detecting the relevant environment variable.
+        if os.getenv('PYTEST_CURRENT_TEST'):
+            self.log.debug('Running in Pytest, setting default `exit_process` to False.')
+            return False
+
+        # Check if running in Scrapy by attempting to import it.
+        with suppress(ImportError):
+            import scrapy  # noqa: F401
+
+            self.log.debug('Running in Scrapy, setting default `exit_process` to False.')
+            return False
+
+        return True
+
 
 Actor = cast(_ActorType, Proxy(_ActorType))
 """The entry point of the SDK, through which all the Actor operations should be done."""

{apify-2.4.0b2 → apify-2.4.0b4}/src/apify/_crypto.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import base64
+import hashlib
+import hmac
+import string
 from typing import Any
 
 from cryptography.exceptions import InvalidTag as InvalidTagException
@@ -153,3 +156,38 @@ def decrypt_input_secrets(private_key: rsa.RSAPrivateKey, input_data: Any) -> An
                 )
 
     return input_data
+
+
+CHARSET = string.digits + string.ascii_letters
+
+
+def encode_base62(num: int) -> str:
+    """Encode the given number to base62."""
+    if num == 0:
+        return CHARSET[0]
+
+    res = ''
+    while num > 0:
+        num, remainder = divmod(num, 62)
+        res = CHARSET[remainder] + res
+    return res
+
+
+@ignore_docs
+def create_hmac_signature(secret_key: str, message: str) -> str:
+    """Generate an HMAC signature and encodes it using Base62. Base62 encoding reduces the signature length.
+
+    HMAC signature is truncated to 30 characters to make it shorter.
+
+    Args:
+        secret_key: Secret key used for signing signatures.
+        message: Message to be signed.
+
+    Returns:
+        Base62 encoded signature.
+    """
+    signature = hmac.new(secret_key.encode('utf-8'), message.encode('utf-8'), hashlib.sha256).hexdigest()[:30]
+
+    decimal_signature = int(signature, 16)
+
+    return encode_base62(decimal_signature)

{apify-2.4.0b2 → apify-2.4.0b4}/src/apify/_platform_event_manager.py

@@ -4,7 +4,7 @@ import asyncio
 from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
 
-import websockets.client
+import websockets.asyncio.client
 from pydantic import BaseModel, Discriminator, Field, TypeAdapter
 from typing_extensions import Self, Unpack, override
 
@@ -143,7 +143,7 @@ class PlatformEventManager(EventManager):
     but instead use it via the `Actor.on()` and `Actor.off()` methods.
     """
 
-    _platform_events_websocket: websockets.client.WebSocketClientProtocol | None = None
+    _platform_events_websocket: websockets.asyncio.client.ClientConnection | None = None
     _process_platform_messages_task: asyncio.Task | None = None
     _send_system_info_interval_task: asyncio.Task | None = None
     _connected_to_platform_websocket: asyncio.Future = asyncio.Future()
@@ -196,7 +196,7 @@ class PlatformEventManager(EventManager):
 
     async def _process_platform_messages(self, ws_url: str) -> None:
         try:
-            async with websockets.client.connect(ws_url) as websocket:
+            async with websockets.asyncio.client.connect(ws_url) as websocket:
                 self._platform_events_websocket = websocket
                 self._connected_to_platform_websocket.set_result(True)
 

{apify-2.4.0b2 → apify-2.4.0b4}/src/apify/apify_storage_client/_key_value_store_client.py

@@ -4,10 +4,13 @@ from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
 from typing_extensions import override
+from yarl import URL
 
 from crawlee.storage_clients._base import KeyValueStoreClient as BaseKeyValueStoreClient
 from crawlee.storage_clients.models import KeyValueStoreListKeysPage, KeyValueStoreMetadata, KeyValueStoreRecord
 
+from apify._crypto import create_hmac_signature
+
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
     from contextlib import AbstractAsyncContextManager
@@ -89,6 +92,18 @@ class KeyValueStoreClient(BaseKeyValueStoreClient):
         Args:
             key: The key for which the URL should be generated.
         """
-        public_api_url = self._api_public_base_url
+        if self._client.resource_id is None:
+            raise ValueError('resource_id cannot be None when generating a public URL')
+
+        public_url = (
+            URL(self._api_public_base_url) / 'v2' / 'key-value-stores' / self._client.resource_id / 'records' / key
+        )
+
+        key_value_store = await self.get()
+
+        if key_value_store is not None and isinstance(key_value_store.model_extra, dict):
+            url_signing_secret_key = key_value_store.model_extra.get('urlSigningSecretKey')
+            if url_signing_secret_key:
+                public_url = public_url.with_query(signature=create_hmac_signature(url_signing_secret_key, key))
 
-        return f'{public_api_url}/v2/key-value-stores/{self._client.resource_id}/records/{key}'
+        return str(public_url)

{apify-2.4.0b2 → apify-2.4.0b4}/tests/integration/test_actor_key_value_store.py

@@ -201,19 +201,28 @@ async def test_generate_public_url_for_kvs_record(
     run_actor: RunActorFunction,
 ) -> None:
     async def main() -> None:
-        from typing import cast
-
-        from apify.apify_storage_client._key_value_store_client import KeyValueStoreClient
+        from apify._crypto import create_hmac_signature
 
         async with Actor:
             public_api_url = Actor.config.api_public_base_url
             default_store_id = Actor.config.default_key_value_store_id
+            record_key = 'public-record-key'
 
             store = await Actor.open_key_value_store()
-            record_url = await cast(KeyValueStoreClient, store._resource_client).get_public_url('dummy')
-            print(record_url)
 
-            assert record_url == f'{public_api_url}/v2/key-value-stores/{default_store_id}/records/dummy'
+            assert isinstance(store.storage_object.model_extra, dict)
+            url_signing_secret_key = store.storage_object.model_extra.get('urlSigningSecretKey')
+            assert url_signing_secret_key is not None
+
+            await store.set_value(record_key, {'exposedData': 'test'}, 'application/json')
+
+            record_url = await store.get_public_url(record_key)
+
+            signature = create_hmac_signature(url_signing_secret_key, record_key)
+            assert (
+                record_url
+                == f'{public_api_url}/v2/key-value-stores/{default_store_id}/records/{record_key}?signature={signature}'
+            )
 
     actor = await make_actor(label='kvs-get-public-url', main_func=main)
     run_result = await run_actor(actor)

{apify-2.4.0b2 → apify-2.4.0b4}/tests/unit/actor/test_actor_lifecycle.py

@@ -9,7 +9,7 @@ from typing import Any, Callable, cast
 from unittest.mock import AsyncMock, Mock
 
 import pytest
-import websockets.server
+import websockets.asyncio.server
 from lazy_object_proxy import Proxy
 
 from apify_shared.consts import ActorEnvVars, ApifyEnvVars
@@ -139,10 +139,10 @@ async def test_actor_handles_migrating_event_correctly(monkeypatch: pytest.Monke
         nonlocal persist_state_events_data
         persist_state_events_data.append(data)
 
-    async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
+    async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
         await websocket.wait_closed()
 
-    async with websockets.server.serve(handler, host='localhost') as ws_server:
+    async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
         port: int = ws_server.sockets[0].getsockname()[1]  # type: ignore[index]
         monkeypatch.setenv(ApifyEnvVars.ACTOR_EVENTS_WS_URL, f'ws://localhost:{port}')
 
@@ -181,7 +181,7 @@ async def test_actor_handles_migrating_event_correctly(monkeypatch: pytest.Monke
             Actor.on(Event.PERSIST_STATE, log_persist_state)
             await asyncio.sleep(2)
 
-            for socket in ws_server.websockets:
+            for socket in ws_server.connections:
                 await socket.send(
                     json.dumps(
                         {

apify-2.4.0b4/tests/unit/actor/test_actor_log.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import contextlib
+import logging
+from typing import TYPE_CHECKING
+
+from apify import Actor
+from apify.log import logger
+
+if TYPE_CHECKING:
+    import pytest
+
+
+async def test_actor_logs_messages_correctly(caplog: pytest.LogCaptureFixture) -> None:
+    caplog.set_level(logging.DEBUG, logger='apify')
+
+    with contextlib.suppress(RuntimeError):
+        async with Actor(configure_logging=False):
+            # Test Actor.log
+            Actor.log.debug('Debug message')
+            Actor.log.info('Info message')
+
+            # Test logger
+            logger.warning('Warning message')
+            logger.error('Error message')
+
+            # Test that exception is logged with the traceback
+            try:
+                raise ValueError('Dummy ValueError')
+            except Exception:
+                Actor.log.exception('Exception message')
+
+            # Test multiline message being indented correctly
+            logger.info('Multi\nline\nlog\nmessage')
+
+            # Test that exception in Actor.main is logged with the traceback
+            raise RuntimeError('Dummy RuntimeError')
+
+    # Updated expected number of log records (an extra record is now captured)
+    assert len(caplog.records) == 14
+
+    # Record 0: Extra Pytest context log
+    assert caplog.records[0].levelno == logging.DEBUG
+    assert caplog.records[0].message.startswith('Running in Pytest')
+
+    # Record 1: Duplicate Pytest context log
+    assert caplog.records[1].levelno == logging.DEBUG
+    assert caplog.records[0].message.startswith('Running in Pytest')
+
+    # Record 2: Initializing Actor...
+    assert caplog.records[2].levelno == logging.INFO
+    assert caplog.records[2].message == 'Initializing Actor...'
+
+    # Record 3: System info
+    assert caplog.records[3].levelno == logging.INFO
+    assert caplog.records[3].message == 'System info'
+
+    # Record 4: Event manager initialized
+    assert caplog.records[4].levelno == logging.DEBUG
+    assert caplog.records[4].message == 'Event manager initialized'
+
+    # Record 5: Charging manager initialized
+    assert caplog.records[5].levelno == logging.DEBUG
+    assert caplog.records[5].message == 'Charging manager initialized'
+
+    # Record 6: Debug message
+    assert caplog.records[6].levelno == logging.DEBUG
+    assert caplog.records[6].message == 'Debug message'
+
+    # Record 7: Info message
+    assert caplog.records[7].levelno == logging.INFO
+    assert caplog.records[7].message == 'Info message'
+
+    # Record 8: Warning message
+    assert caplog.records[8].levelno == logging.WARNING
+    assert caplog.records[8].message == 'Warning message'
+
+    # Record 9: Error message
+    assert caplog.records[9].levelno == logging.ERROR
+    assert caplog.records[9].message == 'Error message'
+
+    # Record 10: Exception message with traceback (ValueError)
+    assert caplog.records[10].levelno == logging.ERROR
+    assert caplog.records[10].message == 'Exception message'
+    assert caplog.records[10].exc_info is not None
+    assert caplog.records[10].exc_info[0] is ValueError
+    assert isinstance(caplog.records[10].exc_info[1], ValueError)
+    assert str(caplog.records[10].exc_info[1]) == 'Dummy ValueError'
+
+    # Record 11: Multiline log message
+    assert caplog.records[11].levelno == logging.INFO
+    assert caplog.records[11].message == 'Multi\nline\nlog\nmessage'
+
+    # Record 12: Actor failed with an exception (RuntimeError)
+    assert caplog.records[12].levelno == logging.ERROR
+    assert caplog.records[12].message == 'Actor failed with an exception'
+    assert caplog.records[12].exc_info is not None
+    assert caplog.records[12].exc_info[0] is RuntimeError
+    assert isinstance(caplog.records[12].exc_info[1], RuntimeError)
+    assert str(caplog.records[12].exc_info[1]) == 'Dummy RuntimeError'
+
+    # Record 13: Exiting Actor
+    assert caplog.records[13].levelno == logging.INFO
+    assert caplog.records[13].message == 'Exiting Actor'

{apify-2.4.0b2 → apify-2.4.0b4}/tests/unit/test_crypto.py

@@ -4,7 +4,15 @@ import base64
 
 import pytest
 
-from apify._crypto import _load_public_key, crypto_random_object_id, load_private_key, private_decrypt, public_encrypt
+from apify._crypto import (
+    _load_public_key,
+    create_hmac_signature,
+    crypto_random_object_id,
+    encode_base62,
+    load_private_key,
+    private_decrypt,
+    public_encrypt,
+)
 
 # NOTE: Uses the same keys as in:
 # https://github.com/apify/apify-shared-js/blob/master/test/crypto.test.ts
@@ -105,3 +113,25 @@ def test_crypto_random_object_id_length_and_charset() -> None:
     long_random_object_id = crypto_random_object_id(1000)
     for char in long_random_object_id:
         assert char in 'abcdefghijklmnopqrstuvwxyzABCEDFGHIJKLMNOPQRSTUVWXYZ0123456789'
+
+
+@pytest.mark.parametrize(('test_input', 'expected'), [(0, '0'), (10, 'a'), (999999999, '15FTGf')])
+def test_encode_base62(test_input: int, expected: str) -> None:
+    assert encode_base62(test_input) == expected
+
+
+# This test ensures compatibility with the JavaScript version of the same method.
+# https://github.com/apify/apify-shared-js/blob/master/packages/utilities/src/hmac.ts
+def test_create_valid_hmac_signature() -> None:
+    # This test uses the same secret key and message as in JS tests.
+    secret_key = 'hmac-secret-key'
+    message = 'hmac-message-to-be-authenticated'
+    assert create_hmac_signature(secret_key, message) == 'pcVagAsudj8dFqdlg7mG'
+
+
+def test_create_same_hmac() -> None:
+    # This test uses the same secret key and message as in JS tests.
+    secret_key = 'hmac-same-secret-key'
+    message = 'hmac-same-message-to-be-authenticated'
+    assert create_hmac_signature(secret_key, message) == 'FYMcmTIm3idXqleF1Sw5'
+    assert create_hmac_signature(secret_key, message) == 'FYMcmTIm3idXqleF1Sw5'

{apify-2.4.0b2 → apify-2.4.0b4}/tests/unit/test_platform_event_manager.py

@@ -9,7 +9,7 @@ from unittest.mock import Mock
 
 import pytest
 import websockets
-import websockets.server
+import websockets.asyncio.server
 
 from apify_shared.consts import ActorEnvVars
 from crawlee.events._types import Event
@@ -133,16 +133,16 @@ async def test_lifecycle_on_platform_without_websocket(monkeypatch: pytest.Monke
 
 
 async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
-    connected_ws_clients: set[websockets.server.WebSocketServerProtocol] = set()
+    connected_ws_clients: set[websockets.asyncio.server.ServerConnection] = set()
 
-    async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
+    async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
         connected_ws_clients.add(websocket)
         try:
             await websocket.wait_closed()
         finally:
             connected_ws_clients.remove(websocket)
 
-    async with websockets.server.serve(handler, host='localhost') as ws_server:
+    async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
         # When you don't specify a port explicitly, the websocket connection is opened on a random free port.
         # We need to find out which port is that.
         port: int = ws_server.sockets[0].getsockname()[1]  # type: ignore[index]
@@ -153,9 +153,9 @@ async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
-    connected_ws_clients: set[websockets.server.WebSocketServerProtocol] = set()
+    connected_ws_clients: set[websockets.asyncio.server.ServerConnection] = set()
 
-    async def handler(websocket: websockets.server.WebSocketServerProtocol) -> None:
+    async def handler(websocket: websockets.asyncio.server.ServerConnection) -> None:
         connected_ws_clients.add(websocket)
         try:
             await websocket.wait_closed()
@@ -169,7 +169,7 @@ async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> No
 
         websockets.broadcast(connected_ws_clients, json.dumps(message))
 
-    async with websockets.server.serve(handler, host='localhost') as ws_server:
+    async with websockets.asyncio.server.serve(handler, host='localhost') as ws_server:
         # When you don't specify a port explicitly, the websocket connection is opened on a random free port.
         # We need to find out which port is that.
         port: int = ws_server.sockets[0].getsockname()[1]  # type: ignore[index]