Finalize tests

Author: Pijukatel

src/apify/_actor.py

@@ -204,9 +204,10 @@ async def __aenter__(self) -> Self:
         if not Actor.is_at_home():
             # Make sure that the input related KVS is initialized to ensure that the input aware client is used
             await self.open_key_value_store()
-
-        # Load non-default aliased storages from configuration
-        await AliasResolver.register_aliases(configuration=self.configuration)
+        else:
+            # Load pre-existing non-default aliased storages from configuration
+            # Supported only on the Apify platform, where those storages are pre-created by the platform.
+            await AliasResolver.register_aliases(configuration=self.configuration)
         return self
 
     async def __aexit__(

src/apify/storage_clients/_apify/_alias_resolving.py

@@ -291,15 +291,12 @@ async def register_aliases(cls, configuration: Configuration) -> None:
                     )._storage_key
                 ] = storage_id
 
-        if configuration.is_at_home:
-            # Bulk update the mapping in the default KVS with the configuration mapping.
-            client = await cls._get_default_kvs_client(configuration=configuration)
-            existing_mapping = ((await client.get_record(cls._ALIAS_MAPPING_KEY)) or {'value': {}}).get('value', {})
-            # Update the existing mapping with the configuration mapping.
-            existing_mapping.update(configuration_mapping)
-            # Store the updated mapping back in the KVS and in memory.
-            await client.set_record(cls._ALIAS_MAPPING_KEY, existing_mapping)
-            cls._alias_map = existing_mapping
-        else:
-            # Update only in-memory mapping when not at home
-            cls._alias_map.update(configuration_mapping)
+        # Bulk update the mapping in the default KVS with the configuration mapping.
+        client = await cls._get_default_kvs_client(configuration=configuration)
+        existing_mapping = ((await client.get_record(cls._ALIAS_MAPPING_KEY)) or {'value': {}}).get('value', {})
+
+        # Update the existing mapping with the configuration mapping.
+        existing_mapping.update(configuration_mapping)
+        # Store the updated mapping back in the KVS and in memory.
+        await client.set_record(cls._ALIAS_MAPPING_KEY, existing_mapping)
+        cls._alias_map.update(existing_mapping)

tests/integration/test_storages.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
 import asyncio
+from typing import cast
 
 import pytest
 
 from crawlee import service_locator
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 
 from apify import Actor, Configuration
+from apify._configuration import ActorStorages
 from apify.storage_clients import ApifyStorageClient, MemoryStorageClient, SmartApifyStorageClient
+from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
 @pytest.mark.parametrize(
@@ -125,3 +128,53 @@ async def test_actor_implicit_storage_init(apify_token: str) -> None:
         assert await Actor.open_dataset() is not await Actor.open_dataset(force_cloud=True)
         assert await Actor.open_key_value_store() is not await Actor.open_key_value_store(force_cloud=True)
         assert await Actor.open_request_queue() is not await Actor.open_request_queue(force_cloud=True)
+
+
+async def test_actor_storages_alias_resolving(apify_token: str) -> None:
+    """Test that `AliasResolver.register_aliases` correctly updates default KVS with Actor storages."""
+
+    # Actor storages
+    datasets = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
+    request_queues = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
+    key_value_stores = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
+
+    # Set up the configuration and storage client for the test
+    configuration = Configuration(
+        default_key_value_store_id='default_kvs_id',
+        token=apify_token,
+        actor_storages=ActorStorages(
+            datasets=datasets, request_queues=request_queues, key_value_stores=key_value_stores
+        ),
+    )
+    storage_client = ApifyStorageClient()
+    service_locator.set_configuration(configuration)
+    service_locator.set_storage_client(storage_client)
+
+    client_cache_key = cast('tuple', storage_client.get_storage_client_cache_key(configuration))[-1]
+    # Add some unrelated pre-existing alias mapping (it should be preserved after registering aliases)
+    pre_existing_mapping = {f'KeyValueStore,pre_existing_alias,{client_cache_key}': 'pre_existing_id'}
+
+    default_kvs = await KeyValueStore.open(configuration=configuration, storage_client=storage_client)
+    await default_kvs.set_value(AliasResolver._ALIAS_MAPPING_KEY, pre_existing_mapping)
+
+    # Construct the expected mapping
+    expected_mapping = {}
+    for storage_type, storage_map in (
+        ('Dataset', datasets),
+        ('KeyValueStore', key_value_stores),
+        ('RequestQueue', request_queues),
+    ):
+        for storage_alias, storage_id in storage_map.items():
+            expected_mapping[
+                ','.join(
+                    (storage_type, '__default__' if storage_alias == 'default' else storage_alias, client_cache_key)
+                )
+            ] = storage_id
+    expected_mapping.update(pre_existing_mapping)
+
+    try:
+        configuration.default_key_value_store_id = default_kvs.id
+        await AliasResolver.register_aliases(configuration=configuration)
+        assert await default_kvs.get_value(AliasResolver._ALIAS_MAPPING_KEY) == expected_mapping
+    finally:
+        await default_kvs.drop()

tests/unit/actor/test_configuration.py

@@ -306,9 +306,9 @@ def test_actor_storage_json_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test that actor_storages_json is parsed from JSON env var."""
     import json
 
-    datasets = {"default":  "default_dataset_id", "custom": "custom_dataset_id"}
-    request_queues = {"default": "default_dataset_id", "custom": "custom_dataset_id"}
-    key_value_stores = {"default": "default_dataset_id", "custom": "custom_dataset_id"}
+    datasets = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
+    request_queues = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
+    key_value_stores = {'default': 'default_dataset_id', 'custom': 'custom_dataset_id'}
 
     actor_storages_json = json.dumps(
         {

tests/unit/storage_clients/test_alias_resolver.py

@@ -1,10 +1,6 @@
 from __future__ import annotations
 
-from crawlee import service_locator
-
-from apify import Actor
-from apify._configuration import Configuration, ActorStorages
-from apify.storage_clients import SmartApifyStorageClient, ApifyStorageClient
+from apify._configuration import Configuration
 from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
@@ -80,30 +76,3 @@ async def test_get_alias_map_returns_in_memory_map() -> None:
     AliasResolver._alias_map = {}
     result = await AliasResolver._get_alias_map(config)
     assert result == {}
-
-
-async def test_register_aliases() -> None:
-    """Test that _get_alias_map loads the map from KVS when at home.
-
-    AliasResolver works locally only """
-
-
-    datasets = {"default":  "default_dataset_id", "custom": "custom_dataset_id"}
-    request_queues = {"default": "default_dataset_id", "custom": "custom_dataset_id"}
-    key_value_stores = {"default": "default_dataset_id", "custom": "custom_dataset_id"}
-
-    config = Configuration(is_at_home=False,
-                           token='test-token',
-                           actor_storages= ActorStorages(
-                               datasets = datasets,
-                               request_queues = request_queues,
-                               key_value_stores = key_value_stores
-                           ),
-    )
-    storage_client = ApifyStorageClient()
-    service_locator.set_storage_client(
-        SmartApifyStorageClient(local_storage_client=storage_client, cloud_storage_client=storage_client)
-    )
-    async with Actor(configuration=config):
-        d = await Actor.open_dataset(alias='default')
-        assert d