Add user_data test

Author: Pijukatel

src/apify/storage_clients/_apify/_request_queue_single_client.py

@@ -197,6 +197,7 @@ async def _get_request(self, id: str) -> Request | None:
         if id in self._requests_cache:
             return self._requests_cache[id]
 
+        # Requests that were not added by this client are not in local cache. Fetch them from platform.
         response = await self._api_client.get_request(id)
 
         if response is None:
@@ -206,13 +207,13 @@ async def _get_request(self, id: str) -> Request | None:
 
         # Updated local caches
         if id in self._requests_in_progress:
-            # Ignore requests that are already in progress, client is already aware of them.
-            self._requests_already_handled.add(id)
+            # No caching of requests that are already in progress, client is already aware of them.
+            pass
         elif request.was_already_handled:
             # Cache only id for already handled requests
             self._requests_already_handled.add(id)
         else:
-            # Cache full request for unhandled requests that are not yet in progress
+            # Cache full request for unhandled requests that are not yet in progress and are not yet handled.
             self._requests_cache[id] = request
         return request
 
@@ -268,16 +269,9 @@ async def _list_head(self) -> None:
             if request.was_already_handled:
                 # Do not cache fully handled requests, we do not need them. Just cache their id.
                 self._requests_already_handled.add(request_id)
-            else:
-                # Only fetch the request if we do not know it yet.
-                if request_id not in self._requests_cache:
-                    complete_request_data = await self._api_client.get_request(request_id)
-                    request = Request.model_validate(complete_request_data)
-                    self._requests_cache[request_id] = request
-
-                # Add new requests to the end of the head, unless already present in head
-                if request_id not in self._head_requests:
-                    self._head_requests.appendleft(request_id)
+            # Add new requests to the end of the head, unless already present in head
+            elif request_id not in self._head_requests:
+                self._head_requests.appendleft(request_id)
 
     async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
         """Mark a request as handled after successful processing.

tests/integration/test_request_queue.py

@@ -1193,6 +1193,7 @@ async def test_request_queue_has_stats(request_queue_apify: RequestQueue) -> Non
 
 
 async def test_long_request(request_queue_apify: RequestQueue) -> None:
+    rq = request_queue_apify
     request = Request.from_url(
         'https://portal.isoss.gov.cz/irj/portal/anonymous/mvrest?path=/eosm-public-offer&officeLabels=%7B%7D&page=1&pageSize=100000&sortColumn=zdatzvsm&sortOrder=-1',
         use_extended_unique_key=True,
@@ -1201,13 +1202,38 @@ async def test_long_request(request_queue_apify: RequestQueue) -> None:
 
     request_id = unique_key_to_request_id(request.unique_key)
 
-    processed_request = await request_queue_apify.add_request(request)
+    processed_request = await rq.add_request(request)
     assert processed_request.id == request_id
 
-    request_obtained = await request_queue_apify.fetch_next_request()
+    request_obtained = await rq.fetch_next_request()
     assert request_obtained is not None
 
-    await request_queue_apify.mark_request_as_handled(request_obtained)
+    await rq.mark_request_as_handled(request_obtained)
 
-    is_finished = await request_queue_apify.is_finished()
+    is_finished = await rq.is_finished()
     assert is_finished
+
+
+async def test_pre_existing_request_with_user_data(
+    request_queue_apify: RequestQueue, apify_client_async: ApifyClientAsync
+) -> None:
+    """Test that pre-existing requests with user data are fully correctly.
+
+    list_head does not return user data, so we need to test that fetching unknown requests is not relying on it."""
+    custom_data = {'key': 'value'}
+
+    rq = request_queue_apify
+    request = Request.from_url(
+        'https://example.com',
+        user_data=custom_data,
+    )
+
+    # Add request by a different producer
+    rq_client = apify_client_async.request_queue(request_queue_id=rq.id)
+    await rq_client.add_request(request.model_dump(by_alias=True))
+
+    # Fetch the request by the client under test
+    request_obtained = await rq.fetch_next_request()
+    assert request_obtained is not None
+    # Test that custom_data is preserved in user_data (custom_data should be subset of obtained user_data)
+    assert custom_data.items() <= request_obtained.user_data.items()