Add tests for RQ

Author: vdusek

src/crawlee/storage_clients/_file_system/_request_queue_client.py

@@ -104,10 +104,10 @@ async def open(
     ) -> FileSystemRequestQueueClient:
         if id:
             raise ValueError(
-                'Opening a dataset by "id" is not supported for file system storage client, use "name" instead.'
+                'Opening a request queue by "id" is not supported for file system storage client, use "name" instead.'
             )
 
-        name = name or configuration.default_dataset_id
+        name = name or configuration.default_request_queue_id
 
         # Check if the client is already cached by name.
         if name in cls._cache_by_name:
@@ -123,7 +123,7 @@ async def open(
         if rq_path.exists():
             # If metadata file is missing, raise an error.
             if not metadata_path.exists():
-                raise ValueError(f'Metadata file not found for RQ "{name}"')
+                raise ValueError(f'Metadata file not found for request queue "{name}"')
 
             file = await asyncio.to_thread(open, metadata_path)
             try:
@@ -133,7 +133,7 @@ async def open(
             try:
                 metadata = RequestQueueMetadata(**file_content)
             except ValidationError as exc:
-                raise ValueError(f'Invalid metadata file for RQ "{name}"') from exc
+                raise ValueError(f'Invalid metadata file for request queue "{name}"') from exc
 
             client = cls(
                 id=metadata.id,
@@ -185,6 +185,8 @@ async def drop(self) -> None:
         if self.metadata.name in self.__class__._cache_by_name:  # noqa: SLF001
             del self.__class__._cache_by_name[self.metadata.name]  # noqa: SLF001
 
+    # TODO: continue
+
     @override
     async def add_batch_of_requests(
         self,

src/crawlee/storages/_key_value_store.py

@@ -99,7 +99,7 @@ async def open(
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
-        # Check if key value store is already cached by id or name
+        # Check if key-value store is already cached by id or name
         if id and id in cls._cache_by_id:
             return cls._cache_by_id[id]
         if name and name in cls._cache_by_name:
@@ -116,7 +116,7 @@ async def open(
 
         kvs = cls(client)
 
-        # Cache the key value store by id and name if available
+        # Cache the key-value store by id and name if available
         if kvs.id:
             cls._cache_by_id[kvs.id] = kvs
         if kvs.name:

src/crawlee/storages/_request_queue.py

@@ -76,7 +76,7 @@ def __init__(self, client: RequestQueueClient) -> None:
         Preferably use the `RequestQueue.open` constructor to create a new instance.
 
         Args:
-            client: An instance of a key-value store client.
+            client: An instance of a request queue client.
         """
         self._client = client
 
@@ -111,7 +111,7 @@ async def open(
         if id and name:
             raise ValueError('Only one of "id" or "name" can be specified, not both.')
 
-        # Check if key value store is already cached by id or name
+        # Check if request queue is already cached by id or name
         if id and id in cls._cache_by_id:
             return cls._cache_by_id[id]
         if name and name in cls._cache_by_name:
@@ -126,7 +126,15 @@ async def open(
             configuration=configuration,
         )
 
-        return cls(client)
+        rq = cls(client)
+
+        # Cache the request queue by id and name if available
+        if rq.id:
+            cls._cache_by_id[rq.id] = rq
+        if rq.name:
+            cls._cache_by_name[rq.name] = rq
+
+        return rq
 
     @override
     async def drop(self) -> None:
@@ -163,27 +171,32 @@ async def add_requests(
         transformed_requests = self._transform_requests(requests)
         wait_time_secs = wait_time_between_batches.total_seconds()
 
-        async def _process_batch(batch: Sequence[Request]) -> None:
-            request_count = len(batch)
-            response = await self._client.add_batch_of_requests(batch, forefront=forefront)
-            logger.debug(f'Added {request_count} requests to the queue, response: {response}')
-
         # Wait for the first batch to be added
         first_batch = transformed_requests[:batch_size]
         if first_batch:
-            await _process_batch(first_batch)
+            await self._process_batch(
+                first_batch,
+                base_retry_wait=wait_time_between_batches,
+                forefront=forefront,
+            )
 
         async def _process_remaining_batches() -> None:
             for i in range(batch_size, len(transformed_requests), batch_size):
                 batch = transformed_requests[i : i + batch_size]
-                await _process_batch(batch)
+                await self._process_batch(
+                    batch,
+                    base_retry_wait=wait_time_between_batches,
+                    forefront=forefront,
+                )
                 if i + batch_size < len(transformed_requests):
                     await asyncio.sleep(wait_time_secs)
 
         # Create and start the task to process remaining batches in the background
         remaining_batches_task = asyncio.create_task(
-            _process_remaining_batches(), name='request_queue_process_remaining_batches_task'
+            _process_remaining_batches(),
+            name='request_queue_process_remaining_batches_task',
         )
+
         self._add_requests_tasks.append(remaining_batches_task)
         remaining_batches_task.add_done_callback(lambda _: self._add_requests_tasks.remove(remaining_batches_task))
 
@@ -195,69 +208,6 @@ async def _process_remaining_batches() -> None:
                 timeout=wait_for_all_requests_to_be_added_timeout,
             )
 
-        # Wait for the first batch to be added
-        first_batch = transformed_requests[:batch_size]
-        if first_batch:
-            await self._process_batch(first_batch, base_retry_wait=wait_time_between_batches)
-
-        async def _process_remaining_batches() -> None:
-            for i in range(batch_size, len(transformed_requests), batch_size):
-                batch = transformed_requests[i : i + batch_size]
-                await self._process_batch(batch, base_retry_wait=wait_time_between_batches)
-                if i + batch_size < len(transformed_requests):
-                    await asyncio.sleep(wait_time_secs)
-
-        # Create and start the task to process remaining batches in the background
-        remaining_batches_task = asyncio.create_task(
-            _process_remaining_batches(), name='request_queue_process_remaining_batches_task'
-        )
-        self._tasks.append(remaining_batches_task)
-        remaining_batches_task.add_done_callback(lambda _: self._tasks.remove(remaining_batches_task))
-
-        # Wait for all tasks to finish if requested
-        if wait_for_all_requests_to_be_added:
-            await wait_for_all_tasks_for_finish(
-                (remaining_batches_task,),
-                logger=logger,
-                timeout=wait_for_all_requests_to_be_added_timeout,
-            )
-
-    async def _process_batch(self, batch: Sequence[Request], base_retry_wait: timedelta, attempt: int = 1) -> None:
-        max_attempts = 5
-        response = await self._resource_client.batch_add_requests(batch)
-
-        if response.unprocessed_requests:
-            logger.debug(f'Following requests were not processed: {response.unprocessed_requests}.')
-            if attempt > max_attempts:
-                logger.warning(
-                    f'Following requests were not processed even after {max_attempts} attempts:\n'
-                    f'{response.unprocessed_requests}'
-                )
-            else:
-                logger.debug('Retry to add requests.')
-                unprocessed_requests_unique_keys = {request.unique_key for request in response.unprocessed_requests}
-                retry_batch = [request for request in batch if request.unique_key in unprocessed_requests_unique_keys]
-                await asyncio.sleep((base_retry_wait * attempt).total_seconds())
-                await self._process_batch(retry_batch, base_retry_wait=base_retry_wait, attempt=attempt + 1)
-
-        request_count = len(batch) - len(response.unprocessed_requests)
-        self._assumed_total_count += request_count
-        if request_count:
-            logger.debug(
-                f'Added {request_count} requests to the queue. Processed requests: {response.processed_requests}'
-            )
-
-    async def get_request(self, request_id: str) -> Request | None:
-        """Retrieve a request from the queue.
-
-        Args:
-            request_id: ID of the request to retrieve.
-
-        Returns:
-            The retrieved request, or `None`, if it does not exist.
-        """
-        # TODO: implement
-
     async def fetch_next_request(self) -> Request | None:
         """Return the next request in the queue to be processed.
 
@@ -346,3 +296,35 @@ async def is_finished(self) -> bool:
             return True
 
         return False
+
+    async def _process_batch(
+        self,
+        batch: Sequence[Request],
+        *,
+        base_retry_wait: timedelta,
+        attempt: int = 1,
+        forefront: bool = False,
+    ) -> None:
+        max_attempts = 5
+        response = await self._client.add_batch_of_requests(batch, forefront=forefront)
+
+        if response.unprocessed_requests:
+            logger.debug(f'Following requests were not processed: {response.unprocessed_requests}.')
+            if attempt > max_attempts:
+                logger.warning(
+                    f'Following requests were not processed even after {max_attempts} attempts:\n'
+                    f'{response.unprocessed_requests}'
+                )
+            else:
+                logger.debug('Retry to add requests.')
+                unprocessed_requests_unique_keys = {request.unique_key for request in response.unprocessed_requests}
+                retry_batch = [request for request in batch if request.unique_key in unprocessed_requests_unique_keys]
+                await asyncio.sleep((base_retry_wait * attempt).total_seconds())
+                await self._process_batch(retry_batch, base_retry_wait=base_retry_wait, attempt=attempt + 1)
+
+        request_count = len(batch) - len(response.unprocessed_requests)
+
+        if request_count:
+            logger.debug(
+                f'Added {request_count} requests to the queue. Processed requests: {response.processed_requests}'
+            )