Update tests

Author: Pijukatel

src/apify/storage_clients/_apify/_utils.py

@@ -38,7 +38,10 @@ def __init__(self, storage_type: _StorageT, alias: str, token: str, api_url: str
         self.token = token
 
     @classmethod
-    def get_additional_cache_key(cls, api_url: str, token: str) -> str:
+    def get_additional_cache_key(cls, api_url: str, token: str, *, encrypted: bool = False) -> str:
+        if not encrypted:
+            return cls.ADDITIONAL_CACHE_KEY_SEPARATOR.join([api_url, token])
+
         encryption_key = Configuration.get_global_configuration().token
         if encryption_key is not None:
             encrypted_token = cls._create_fernet(encryption_key).encrypt(token.encode()).decode()
@@ -47,7 +50,7 @@ def get_additional_cache_key(cls, api_url: str, token: str) -> str:
 
     @property
     def additional_cache_key(self) -> str:
-        return self.get_additional_cache_key(self.api_url, self.token)
+        return self.get_additional_cache_key(self.api_url, self.token, encrypted=False)
 
     @classmethod
     def from_exported_string(cls, alias_as_string: str) -> _Alias:
@@ -67,8 +70,14 @@ def from_exported_string(cls, alias_as_string: str) -> _Alias:
 
         return cls(storage_type=storage_map[storage_class_name], alias=alias, api_url=api_url, token=token)
 
-    def __str__(self) -> str:
-        return self.ALIAS_SEPARATOR.join([self.storage_type.__name__, self.alias, self.additional_cache_key])
+    def get_storage_key(self) -> str:
+        return self.ALIAS_SEPARATOR.join(
+            [
+                self.storage_type.__name__,
+                self.alias,
+                self.get_additional_cache_key(api_url=self.api_url, token=self.token, encrypted=True),
+            ]
+        )
 
     @staticmethod
     def _create_fernet(token: str) -> Fernet:
@@ -91,10 +100,11 @@ async def store_mapping_to_apify_kvs(self, storage_id: str) -> None:
                 record = record['value']
 
             # Update or create the record with the new alias mapping
+
             if isinstance(record, dict):
-                record[str(self)] = storage_id
+                record[self.get_storage_key()] = storage_id
             else:
-                record = {str(self): storage_id}
+                record = {self.get_storage_key(): storage_id}
 
             # Store the mapping back in the KVS.
             await default_kvs_client.set_record(_ALIAS_MAPPING_KEY, record)

tests/unit/actor/test_apify_storage.py

@@ -1,8 +1,10 @@
+from datetime import datetime, timezone
 from unittest import mock
 from unittest.mock import AsyncMock
 
 import pytest
 
+from crawlee.storage_clients.models import StorageMetadata
 from crawlee.storages import Dataset, KeyValueStore, RequestQueue
 from crawlee.storages._base import Storage
 
@@ -23,7 +25,12 @@ async def test_get_additional_cache_key(
     storage: Storage, _storage_client: ApifyDatasetClient | ApifyKeyValueStoreClient | ApifyRequestQueueClient
 ) -> None:
     """Test that Storages based on `ApifyStorageClient` include `token` and `api_base_url` in additional cache key."""
-    storage_names = iter(['1', '2', '3', '1', '3'])
+
+    def create_metadata(id: str) -> StorageMetadata:
+        now = datetime.now(tz=timezone.utc)
+        return StorageMetadata(id=id, name=None, accessed_at=now, created_at=now, modified_at=now)
+
+    storage_ids = iter(['1', '2', '3', '1', '3'])
 
     apify_storage_client = ApifyStorageClient()
 
@@ -34,8 +41,9 @@ async def test_get_additional_cache_key(
     config_4 = Configuration(token='a')
     config_5 = Configuration(token='a', api_base_url='https://super_custom_api.com')
 
-    mocked_open = AsyncMock(spec=_storage_client.open)
-    mocked_open.get_metadata = AsyncMock(storage_names)
+    mocked_client = AsyncMock(spec=type[_storage_client])
+    mocked_client.get_metadata = AsyncMock(side_effect=lambda: create_metadata(next(storage_ids)))
+    mocked_open = AsyncMock(spec=_storage_client.open, return_value=mocked_client)
 
     with mock.patch.object(_storage_client, 'open', mocked_open):
         storage_1 = await storage.open(storage_client=apify_storage_client, configuration=config_1)