Add uni tests, WIP

TODO - how should it behave locally?

Author: Pijukatel

src/apify/_actor.py

@@ -204,9 +204,9 @@ 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()
-        else:
-            # Load non-default aliased storages from configuration
-            await AliasResolver.register_aliases(configuration=self.configuration)
+
+        # Load non-default aliased storages from configuration
+        await AliasResolver.register_aliases(configuration=self.configuration)
         return self
 
     async def __aexit__(

src/apify/_configuration.py

@@ -44,11 +44,12 @@ class ActorStorages:
     request_queues: dict[str, str]
 
 
-def _load_storage_keys(data: None | str | dict) -> ActorStorages | None:
+def _load_storage_keys(data: None | str | dict | ActorStorages) -> ActorStorages | None:
     """Load storage keys from environment."""
     if data is None:
         return None
-
+    if isinstance(data, ActorStorages):
+        return data
     storage_mapping = data if isinstance(data, dict) else json.loads(data)
     return ActorStorages(
         key_value_stores=storage_mapping.get('keyValueStores', {}),

src/apify/storage_clients/_apify/_alias_resolving.py

@@ -291,6 +291,15 @@ async def register_aliases(cls, configuration: Configuration) -> None:
                     )._storage_key
                 ] = storage_id
 
-        client = await cls._get_default_kvs_client(configuration=configuration)
-        existing_mapping = ((await client.get_record(cls._ALIAS_MAPPING_KEY)) or {'value': {}}).get('value', {})
-        await client.set_record(cls._ALIAS_MAPPING_KEY, {**existing_mapping, **configuration_mapping})
+        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)

tests/unit/actor/test_configuration.py

@@ -300,3 +300,26 @@ def test_actor_pricing_info_from_json_env_var(monkeypatch: pytest.MonkeyPatch) -
     config = ApifyConfiguration()
     assert config.actor_pricing_info is not None
     assert config.actor_pricing_info.pricing_model == 'PAY_PER_EVENT'
+
+
+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"}
+
+    actor_storages_json = json.dumps(
+        {
+            'datasets': datasets,
+            'requestQueues': request_queues,
+            'keyValueStores': key_value_stores,
+        }
+    )
+    monkeypatch.setenv('ACTOR_STORAGES_JSON', actor_storages_json)
+    config = ApifyConfiguration()
+    assert config.actor_storages
+    assert config.actor_storages.datasets == datasets
+    assert config.actor_storages.request_queues == request_queues
+    assert config.actor_storages.key_value_stores == key_value_stores

tests/unit/storage_clients/test_alias_resolver.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from apify._configuration import Configuration
+from crawlee import service_locator
+
+from apify import Actor
+from apify._configuration import Configuration, ActorStorages
+from apify.storage_clients import SmartApifyStorageClient, ApifyStorageClient
 from apify.storage_clients._apify._alias_resolving import AliasResolver
 
 
@@ -76,3 +80,30 @@ 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