apify--crawlee/test/core/request_manager_tandem.test.ts

import { log, Request, RequestList, RequestManagerTandem, RequestQueue } from '@crawlee/core';
import { afterAll, beforeAll, beforeEach, describe, expect, test, vi } from 'vitest';

import { MemoryStorageEmulator } from '../shared/MemoryStorageEmulator';

describe('RequestManagerTandem', () => {
    let logLevel: number;
    const emulator = new MemoryStorageEmulator();

    beforeAll(() => {
        logLevel = log.getLevel();
        log.setLevel(log.LEVELS.ERROR);
    });

    beforeEach(async () => {
        await emulator.init();
        vi.restoreAllMocks();
    });

    afterAll(async () => {
        log.setLevel(logLevel);
        await emulator.destroy();
    });

    test('fetchNextRequest transfers from list to queue when queue is empty', async () => {
        // Create sources with 3 URLs
        const requestList = await RequestList.open(null, [
            { url: 'https://example.com/1' },
            { url: 'https://example.com/2' },
            { url: 'https://example.com/3' },
        ]);
        const requestQueue = await RequestQueue.open();

        // Mock the addRequest function of requestQueue to verify it's called
        const addRequestSpy = vi.spyOn(requestQueue, 'addRequest');

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // First fetch should trigger transfer from list to queue
        const request1 = await tandem.fetchNextRequest();

        // Verify the request was transferred from list to queue
        expect(addRequestSpy).toHaveBeenCalled();
        expect(request1).not.toBeNull();
        expect(request1?.url).toBe('https://example.com/1');

        // Fetch more requests to ensure they all come from the queue
        const request2 = await tandem.fetchNextRequest();
        expect(request2).not.toBeNull();
        expect(request2?.url).toBe('https://example.com/2');

        const request3 = await tandem.fetchNextRequest();
        expect(request3).not.toBeNull();
        expect(request3?.url).toBe('https://example.com/3');

        // No more requests should be available
        const request4 = await tandem.fetchNextRequest();
        expect(request4).toBeNull();
    });

    test('markRequestHandled properly marks request as handled in the queue', async () => {
        const requestList = await RequestList.open(null, [{ url: 'https://example.com/1' }]);
        const requestQueue = await RequestQueue.open();

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Mock markRequestHandled in requestQueue
        const markHandledSpy = vi.spyOn(requestQueue, 'markRequestHandled');

        // First fetch a request
        const request = await tandem.fetchNextRequest();
        expect(request).not.toBeNull();

        // Mark it as handled
        await tandem.markRequestHandled(request!);

        // Verify the queue's markRequestHandled was called
        expect(markHandledSpy).toHaveBeenCalledWith(request);
    });

    test('reclaimRequest properly reclaims request in the queue', async () => {
        const requestList = await RequestList.open(null, [{ url: 'https://example.com/1' }]);
        const requestQueue = await RequestQueue.open();

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Mock reclaimRequest in requestQueue
        const reclaimSpy = vi.spyOn(requestQueue, 'reclaimRequest');

        // First fetch a request
        const request = await tandem.fetchNextRequest();
        expect(request).not.toBeNull();

        // Reclaim the request
        await tandem.reclaimRequest(request!);

        // Verify the queue's reclaimRequest was called
        expect(reclaimSpy).toHaveBeenCalledWith(request, undefined);
    });

    test('handledCount returns the queue handledCount', async () => {
        const requestList = await RequestList.open(null, [
            { url: 'https://example.com/1' },
            { url: 'https://example.com/2' },
        ]);
        const requestQueue = await RequestQueue.open();

        // Mock handledCount methods to return fixed values
        vi.spyOn(requestList, 'handledCount').mockReturnValue(3);
        vi.spyOn(requestQueue, 'handledCount').mockResolvedValue(2);

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Only the request queue counts
        await expect(tandem.handledCount()).resolves.toBe(2);
    });

    test('isFinished returns true only when both list and queue are finished', async () => {
        const requestList = await RequestList.open(null, [{ url: 'https://example.com/1' }]);
        const requestQueue = await RequestQueue.open();

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Mock the isFinished methods
        vi.spyOn(requestList, 'isFinished').mockResolvedValue(false);
        vi.spyOn(requestQueue, 'isFinished').mockResolvedValue(false);

        // Neither is finished, so tandem should not be finished
        expect(await tandem.isFinished()).toBe(false);

        // Only list is finished
        vi.spyOn(requestList, 'isFinished').mockResolvedValue(true);
        vi.spyOn(requestQueue, 'isFinished').mockResolvedValue(false);
        expect(await tandem.isFinished()).toBe(false);

        // Only queue is finished
        vi.spyOn(requestList, 'isFinished').mockResolvedValue(false);
        vi.spyOn(requestQueue, 'isFinished').mockResolvedValue(true);
        expect(await tandem.isFinished()).toBe(false);

        // Both are finished
        vi.spyOn(requestList, 'isFinished').mockResolvedValue(true);
        vi.spyOn(requestQueue, 'isFinished').mockResolvedValue(true);
        expect(await tandem.isFinished()).toBe(true);
    });

    test('isEmpty returns true only when both list and queue are empty', async () => {
        const requestList = await RequestList.open(null, [{ url: 'https://example.com/1' }]);
        const requestQueue = await RequestQueue.open();

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Mock the isEmpty methods
        vi.spyOn(requestList, 'isEmpty').mockResolvedValue(false);
        vi.spyOn(requestQueue, 'isEmpty').mockResolvedValue(false);

        // Neither is empty, so tandem should not be empty
        expect(await tandem.isEmpty()).toBe(false);

        // Only list is empty
        vi.spyOn(requestList, 'isEmpty').mockResolvedValue(true);
        vi.spyOn(requestQueue, 'isEmpty').mockResolvedValue(false);
        expect(await tandem.isEmpty()).toBe(false);

        // Only queue is empty
        vi.spyOn(requestList, 'isEmpty').mockResolvedValue(false);
        vi.spyOn(requestQueue, 'isEmpty').mockResolvedValue(true);
        expect(await tandem.isEmpty()).toBe(false);

        // Both are empty
        vi.spyOn(requestList, 'isEmpty').mockResolvedValue(true);
        vi.spyOn(requestQueue, 'isEmpty').mockResolvedValue(true);
        expect(await tandem.isEmpty()).toBe(true);
    });

    test('handles failed batch transfer appropriately', async () => {
        const requestList = await RequestList.open(null, [
            { url: 'https://example.com/1' },
            { url: 'https://example.com/2' },
        ]);
        const requestQueue = await RequestQueue.open();

        // Mock the queue's addRequests to simulate failure
        vi.spyOn(requestQueue, 'addRequest').mockRejectedValue(new Error('Batch add failed'));

        // Mock the reclaimRequest method to verify it's called
        const reclaimSpy = vi.spyOn(requestList, 'reclaimRequest');

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Attempt to fetch which should trigger the batch transfer
        await tandem.fetchNextRequest();

        // Verify that reclaimRequest was called to reclaim the failed requests
        expect(reclaimSpy).toHaveBeenCalled();
    });

    test('added requests are forwarded to the underlying RequestQueue', async () => {
        const requestList = await RequestList.open(null, []);
        const requestQueue = await RequestQueue.open();

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Mock the addRequest method of the queue
        const addRequestSpy = vi.spyOn(requestQueue, 'addRequestsBatched');

        // Add a new request directly through tandem
        const request = new Request({ url: 'https://example.com/new' });
        await tandem.addRequestsBatched([request]);

        // Verify the request was forwarded to the queue
        expect(addRequestSpy).toHaveBeenCalledWith([request], undefined);
    });

    test('async iterator iterates through all requests', async () => {
        const requestList = await RequestList.open(null, [
            { url: 'https://example.com/1' },
            { url: 'https://example.com/2' },
        ]);

        const requestQueue = await RequestQueue.open();
        await requestQueue.addRequest({ url: 'https://example.com/3' });

        const tandem = new RequestManagerTandem(requestList, requestQueue);

        // Iterate through all requests
        const urls: string[] = [];
        for await (const request of tandem) {
            urls.push(request.url);
        }

        // Verify we got both URLs
        expect(urls).toEqual(['https://example.com/1', 'https://example.com/2', 'https://example.com/3']);
    });
});