Fixes

Author: vdusek

src/apify/storage_clients/_apify/_dataset_client.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import warnings
+from datetime import datetime
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 
@@ -65,7 +66,18 @@ def __init__(
     @override
     async def get_metadata(self) -> DatasetMetadata:
         metadata = await self._api_client.get()
-        return DatasetMetadata.model_validate(metadata)
+
+        if metadata is None:
+            raise ValueError('Failed to retrieve dataset metadata.')
+
+        return DatasetMetadata(
+            id=metadata.id,
+            name=metadata.name,
+            created_at=datetime.fromisoformat(metadata.created_at.replace('Z', '+00:00')),
+            modified_at=datetime.fromisoformat(metadata.modified_at.replace('Z', '+00:00')),
+            accessed_at=datetime.fromisoformat(metadata.accessed_at.replace('Z', '+00:00')),
+            item_count=int(metadata.item_count),
+        )
 
     @classmethod
     async def open(

src/apify/storage_clients/_apify/_key_value_store_client.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import warnings
+from datetime import datetime
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 
@@ -11,7 +12,7 @@
 from crawlee.storage_clients.models import KeyValueStoreRecord, KeyValueStoreRecordMetadata
 
 from ._api_client_creation import create_storage_api_client
-from ._models import ApifyKeyValueStoreMetadata, KeyValueStoreListKeysPage
+from ._models import ApifyKeyValueStoreMetadata
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -54,7 +55,18 @@ def __init__(
     @override
     async def get_metadata(self) -> ApifyKeyValueStoreMetadata:
         metadata = await self._api_client.get()
-        return ApifyKeyValueStoreMetadata.model_validate(metadata)
+
+        if metadata is None:
+            raise ValueError('Failed to retrieve dataset metadata.')
+
+        return ApifyKeyValueStoreMetadata(
+            id=metadata.id,
+            name=metadata.name,
+            created_at=datetime.fromisoformat(metadata.created_at.replace('Z', '+00:00')),
+            modified_at=datetime.fromisoformat(metadata.modified_at.replace('Z', '+00:00')),
+            accessed_at=datetime.fromisoformat(metadata.accessed_at.replace('Z', '+00:00')),
+            url_signing_secret_key=metadata.url_signing_secret_key,
+        )
 
     @classmethod
     async def open(
@@ -143,14 +155,13 @@ async def iterate_keys(
         count = 0
 
         while True:
-            response = await self._api_client.list_keys(exclusive_start_key=exclusive_start_key)
-            list_key_page = KeyValueStoreListKeysPage.model_validate(response)
+            list_key_page = await self._api_client.list_keys(exclusive_start_key=exclusive_start_key)
 
             for item in list_key_page.items:
                 # Convert KeyValueStoreKeyInfo to KeyValueStoreRecordMetadata
                 record_metadata = KeyValueStoreRecordMetadata(
                     key=item.key,
-                    size=item.size,
+                    size=int(item.size),
                     content_type='application/octet-stream',  # Content type not available from list_keys
                 )
                 yield record_metadata

src/apify/storage_clients/_apify/_request_queue_client.py

@@ -78,31 +78,34 @@ async def get_metadata(self) -> ApifyRequestQueueMetadata:
         Returns:
             Request queue metadata with accurate counts and timestamps, combining API data with local estimates.
         """
-        response = await self._api_client.get()
+        metadata = await self._api_client.get()
 
-        if response is None:
+        if metadata is None:
             raise ValueError('Failed to fetch request queue metadata from the API.')
 
-        total_request_count = int(response.total_request_count)
-        handled_request_count = int(response.handled_request_count)
-        pending_request_count = int(response.pending_request_count)
-        created_at = datetime.fromisoformat(response.created_at)
-        modified_at = datetime.fromisoformat(response.modified_at)
-        accessed_at = datetime.fromisoformat(response.accessed_at)
+        total_request_count = int(metadata.total_request_count)
+        handled_request_count = int(metadata.handled_request_count)
+        pending_request_count = int(metadata.pending_request_count)
+        created_at = datetime.fromisoformat(metadata.created_at.replace('Z', '+00:00'))
+        modified_at = datetime.fromisoformat(metadata.modified_at.replace('Z', '+00:00'))
+        accessed_at = datetime.fromisoformat(metadata.accessed_at.replace('Z', '+00:00'))
 
         # Enhance API response with local estimations to account for propagation delays (API data can be delayed
         # by a few seconds, while local estimates are immediately accurate).
         return ApifyRequestQueueMetadata(
-            id=response.id,
-            name=response.name,
+            id=metadata.id,
+            name=metadata.name,
             total_request_count=max(total_request_count, self._implementation.metadata.total_request_count),
             handled_request_count=max(handled_request_count, self._implementation.metadata.handled_request_count),
             pending_request_count=pending_request_count,
             created_at=min(created_at, self._implementation.metadata.created_at),
             modified_at=max(modified_at, self._implementation.metadata.modified_at),
             accessed_at=max(accessed_at, self._implementation.metadata.accessed_at),
-            had_multiple_clients=response.had_multiple_clients or self._implementation.metadata.had_multiple_clients,
-            stats=RequestQueueStats.model_validate(response.stats, by_alias=True),  # ty: ignore[possibly-missing-attribute]
+            had_multiple_clients=metadata.had_multiple_clients or self._implementation.metadata.had_multiple_clients,
+            stats=RequestQueueStats.model_validate(
+                metadata.stats.model_dump(by_alias=True) if metadata.stats else {},
+                by_alias=True,
+            ),
         )
 
     @classmethod
@@ -151,7 +154,7 @@ async def open(
         raw_metadata = await api_client.get()
         if raw_metadata is None:
             raise ValueError('Failed to retrieve request queue metadata from the API.')
-        metadata = ApifyRequestQueueMetadata.model_validate(raw_metadata)
+        metadata = ApifyRequestQueueMetadata.model_validate(raw_metadata.model_dump(by_alias=True))
 
         return cls(
             api_client=api_client,

src/apify/storage_clients/_apify/_request_queue_shared_client.py

@@ -121,18 +121,17 @@ async def add_batch_of_requests(
 
         if new_requests:
             # Prepare requests for API by converting to dictionaries.
-            requests_dict = [
-                request.model_dump(
-                    by_alias=True,
-                )
-                for request in new_requests
-            ]
+            requests_dict = [request.model_dump(by_alias=True) for request in new_requests]
 
             # Send requests to API.
-            api_response = AddRequestsResponse.model_validate(
-                await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront)
+            batch_response = await self._api_client.batch_add_requests(
+                requests=requests_dict,
+                forefront=forefront,
             )
 
+            batch_response_dict = batch_response.model_dump(by_alias=True)
+            api_response = AddRequestsResponse.model_validate(batch_response_dict)
+
             # Add the locally known already present processed requests based on the local cache.
             api_response.processed_requests.extend(already_present_requests)
 
@@ -312,7 +311,8 @@ async def _get_request_by_id(self, request_id: str) -> Request | None:
         if response is None:
             return None
 
-        return Request.model_validate(response)
+        response_dict = response.model_dump(by_alias=True)
+        return Request.model_validate(response_dict)
 
     async def _ensure_head_is_non_empty(self) -> None:
         """Ensure that the queue head has requests if they are available in the queue."""
@@ -442,7 +442,7 @@ async def _list_head(
         self.metadata.had_multiple_clients = list_and_lost_data.had_multiple_clients
 
         for request_data in list_and_lost_data.items:
-            request = Request.model_validate(request_data)
+            request = Request.model_validate(request_data.model_dump(by_alias=True))
             request_id = request_data.id
 
             # Skip requests without ID or unique key
@@ -473,7 +473,8 @@ async def _list_head(
             # After adding new requests to the forefront, any existing leftover locked request is kept in the end.
             self._queue_head.append(leftover_id)
 
-        return RequestQueueHead.model_validate(list_and_lost_data)
+        list_and_lost_dict = list_and_lost_data.model_dump(by_alias=True)
+        return RequestQueueHead.model_validate(list_and_lost_dict)
 
     def _cache_request(
         self,

src/apify/storage_clients/_apify/_request_queue_single_client.py

@@ -147,22 +147,20 @@ async def add_batch_of_requests(
 
         if new_requests:
             # Prepare requests for API by converting to dictionaries.
-            requests_dict = [
-                request.model_dump(
-                    by_alias=True,
-                )
-                for request in new_requests
-            ]
+            requests_dict = [request.model_dump(by_alias=True) for request in new_requests]
 
             # Send requests to API.
-            api_response = AddRequestsResponse.model_validate(
-                await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront)
-            )
+            batch_response = await self._api_client.batch_add_requests(requests=requests_dict, forefront=forefront)
+            batch_response_dict = batch_response.model_dump(by_alias=True)
+            api_response = AddRequestsResponse.model_validate(batch_response_dict)
+
             # Add the locally known already present processed requests based on the local cache.
             api_response.processed_requests.extend(already_present_requests)
+
             # Remove unprocessed requests from the cache
             for unprocessed_request in api_response.unprocessed_requests:
-                self._requests_cache.pop(unique_key_to_request_id(unprocessed_request.unique_key), None)
+                request_id = unique_key_to_request_id(unprocessed_request.unique_key)
+                self._requests_cache.pop(request_id, None)
 
         else:
             api_response = AddRequestsResponse(
@@ -292,12 +290,12 @@ async def _list_head(self) -> None:
         # Should warn once? This might be outside expected context if the other consumers consumes at the same time
 
         if response.queue_modified_at:
-            modified_at = datetime.fromisoformat(response.queue_modified_at)
+            modified_at = datetime.fromisoformat(response.queue_modified_at.replace('Z', '+00:00'))
             self.metadata.modified_at = max(self.metadata.modified_at, modified_at)
 
         # Update the cached data
         for request_data in response.items:
-            request = Request.model_validate(request_data)
+            request = Request.model_validate(request_data.model_dump(by_alias=True))
             request_id = request_data.id
 
             if request_id in self._requests_in_progress:
@@ -329,7 +327,8 @@ async def _get_request_by_id(self, id: str) -> Request | None:
         if response is None:
             return None
 
-        request = Request.model_validate(response)
+        response_dict = response.model_dump(by_alias=True)
+        request = Request.model_validate(response_dict)
 
         # Updated local caches
         if id in self._requests_in_progress:
@@ -380,7 +379,7 @@ async def _init_caches(self) -> None:
         """
         response = await self._api_client.list_requests(limit=10_000)
         for request_data in response.items:
-            request = Request.model_validate(request_data)
+            request = Request.model_validate(request_data.model_dump(by_alias=True))
             request_id = request_data.id
 
             if request.was_already_handled:

tests/integration/actor/test_actor_request_queue.py

@@ -98,7 +98,12 @@ async def main() -> None:
             stats_after = _rq.stats
             Actor.log.info(stats_after)
 
-            assert (stats_after['writeCount'] - stats_before['writeCount']) == 1
+            assert stats_after is not None
+            assert stats_after.write_count is not None
+            assert stats_before is not None
+            assert stats_before.write_count is not None
+
+            assert (stats_after.write_count - stats_before.write_count) == 1
 
     actor = await make_actor(label='rq-deduplication', main_func=main)
     run_result = await run_actor(actor)
@@ -147,7 +152,12 @@ async def main() -> None:
             stats_after = _rq.stats
             Actor.log.info(stats_after)
 
-            assert (stats_after['writeCount'] - stats_before['writeCount']) == 2
+            assert stats_after is not None
+            assert stats_after.write_count is not None
+            assert stats_before is not None
+            assert stats_before.write_count is not None
+
+            assert (stats_after.write_count - stats_before.write_count) == 2
 
     actor = await make_actor(label='rq-deduplication', main_func=main)
     run_result = await run_actor(actor)
@@ -193,6 +203,9 @@ async def main() -> None:
             stats_before = _rq.stats
             Actor.log.info(stats_before)
 
+            assert stats_before is not None
+            assert stats_before.write_count is not None
+
             # Add batches of some new and some already present requests in workers
             async def add_requests_worker() -> None:
                 await rq.add_requests(requests[: next(batch_size)])
@@ -207,7 +220,10 @@ async def add_requests_worker() -> None:
             stats_after = _rq.stats
             Actor.log.info(stats_after)
 
-            assert (stats_after['writeCount'] - stats_before['writeCount']) == len(requests)
+            assert stats_after is not None
+            assert stats_after.write_count is not None
+
+            assert (stats_after.write_count - stats_before.write_count) == len(requests)
 
     actor = await make_actor(label='rq-parallel-deduplication', main_func=main)
     run_result = await run_actor(actor)

tests/integration/apify_api/test_request_queue.py

@@ -930,7 +930,7 @@ async def test_request_queue_had_multiple_clients(
     # Check that it is correctly in the API
     api_response = await api_client.get()
     assert api_response
-    assert api_response['hadMultipleClients'] is True
+    assert api_response.had_multiple_clients is True
 
 
 async def test_request_queue_not_had_multiple_clients(
@@ -949,7 +949,7 @@ async def test_request_queue_not_had_multiple_clients(
     api_client = apify_client_async.request_queue(request_queue_id=rq.id)
     api_response = await api_client.get()
     assert api_response
-    assert api_response['hadMultipleClients'] is False
+    assert api_response.had_multiple_clients is False
 
 
 async def test_request_queue_simple_and_full_at_the_same_time(
@@ -1165,6 +1165,9 @@ async def test_request_queue_deduplication_unprocessed_requests(
     stats_before = _rq.stats
     Actor.log.info(stats_before)
 
+    assert stats_before is not None
+    assert stats_before.write_count is not None
+
     def return_unprocessed_requests(requests: list[dict], *_: Any, **__: Any) -> dict[str, list[dict]]:
         """Simulate API returning unprocessed requests."""
         return {
@@ -1176,7 +1179,7 @@ def return_unprocessed_requests(requests: list[dict], *_: Any, **__: Any) -> dic
         }
 
     with mock.patch(
-        'apify_client.clients.resource_clients.request_queue.RequestQueueClientAsync.batch_add_requests',
+        'apify_client._resource_clients.request_queue.RequestQueueClientAsync.batch_add_requests',
         side_effect=return_unprocessed_requests,
     ):
         # Simulate failed API call for adding requests. Request was not processed and should not be cached.
@@ -1191,4 +1194,7 @@ def return_unprocessed_requests(requests: list[dict], *_: Any, **__: Any) -> dic
     stats_after = _rq.stats
     Actor.log.info(stats_after)
 
-    assert (stats_after['writeCount'] - stats_before['writeCount']) == 1
+    assert stats_after is not None
+    assert stats_after.write_count is not None
+
+    assert (stats_after.write_count - stats_before.write_count) == 1

tests/unit/conftest.py

@@ -62,6 +62,7 @@ def prepare_test_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Callabl
     def _prepare_test_env() -> None:
         if hasattr(apify._actor.Actor, '__wrapped__'):
             delattr(apify._actor.Actor, '__wrapped__')
+
         apify._actor.Actor._is_initialized = False
 
         # Set the environment variable for the local storage directory to the temporary path.