Fix BrowserCrawler tests

janbuchar · janbuchar · commit 3a530b883c9c · 2025-10-15T15:28:21.000+02:00
diff --git a/packages/basic-crawler/src/internals/basic-crawler.ts b/packages/basic-crawler/src/internals/basic-crawler.ts
@@ -1336,6 +1336,8 @@ export class BasicCrawler<
         const statisticsId = request.id || request.uniqueKey;
         this.stats.startJob(statisticsId);
 
+        const deferredCleanup: (() => Promise<unknown>)[] = [];
+
         const crawlingContext: CrawlingContext = {
             id: cryptoRandomObjectId(10),
             log: this.log,
@@ -1357,6 +1359,9 @@ export class BasicCrawler<
             useState: this.useState.bind(this),
             sendRequest: createSendRequest(this.httpClient, request!, session, () => crawlingContext.proxyInfo?.url),
             getKeyValueStore: async (idOrName?: string) => KeyValueStore.open(idOrName, { config: this.config }),
+            registerDeferredCleanup: (cleanup) => {
+                deferredCleanup.push(cleanup);
+            },
         };
 
         let isRequestLocked = true;
@@ -1419,6 +1424,8 @@ export class BasicCrawler<
             // decrease the session score if the request fails (but the error handler did not throw)
             crawlingContext.session?.markBad();
         } finally {
+            await Promise.all(deferredCleanup.map((cleanup) => cleanup()));
+
             // Safety net - release the lock if nobody managed to do it before
             if (isRequestLocked && source instanceof RequestProvider) {
                 try {
diff --git a/packages/browser-crawler/src/internals/browser-crawler.ts b/packages/browser-crawler/src/internals/browser-crawler.ts
@@ -1,6 +1,7 @@
 import type {
     Awaitable,
     BasicCrawlerOptions,
+    BasicCrawlingContext,
     CrawlingContext,
     Dictionary,
     EnqueueLinksOptions,
@@ -354,7 +355,8 @@ export abstract class BrowserCrawler<
                 contextPipelineBuilder: () =>
                     contextPipelineBuilder()
                         .compose({ action: this.performNavigation.bind(this) })
-                        .compose({ action: this.handleBlockedRequestByContent.bind(this) }),
+                        .compose({ action: this.handleBlockedRequestByContent.bind(this) })
+                        .compose({ action: this.restoreRequestState.bind(this) }),
                 contextPipelineEnhancer: options.contextPipelineEnhancer,
             },
             config,
@@ -405,10 +407,15 @@ export abstract class BrowserCrawler<
     > {
         return ContextPipeline.create<CrawlingContext>().compose({
             action: this.preparePage.bind(this),
-            cleanup: async (context: { page: Page }) => {
-                await context.page
-                    .close()
-                    .catch((error: Error) => this.log.debug('Error while closing page', { error }));
+            cleanup: async (context: {
+                page: Page;
+                registerDeferredCleanup: BasicCrawlingContext['registerDeferredCleanup'];
+            }) => {
+                context.registerDeferredCleanup(async () => {
+                    await context.page
+                        .close()
+                        .catch((error: Error) => this.log.debug('Error while closing page', { error }));
+                });
             },
         });
     }
@@ -592,6 +599,11 @@ export abstract class BrowserCrawler<
         return {};
     }
 
+    private async restoreRequestState(crawlingContext: CrawlingContext) {
+        crawlingContext.request.state = RequestState.REQUEST_HANDLER;
+        return {};
+    }
+
     protected async _applyCookies(
         { session, request, page, browserController }: BrowserCrawlingContext,
         preHooksCookies: string,
diff --git a/packages/core/src/crawlers/crawler_commons.ts b/packages/core/src/crawlers/crawler_commons.ts
@@ -157,6 +157,11 @@ export interface CrawlingContext<UserData extends Dictionary = Dictionary> exten
      * ```
      */
     sendRequest<Response = string>(overrideOptions?: Partial<OptionsInit>): Promise<GotResponse<Response>>;
+
+    /**
+     * Register a function to be called at the very end of the request handling process. This is useful for resources that should be accessible to error handlers, for instance.
+     */
+    registerDeferredCleanup(cleanup: () => Promise<unknown>): void;
 }
 
 /**
diff --git a/packages/playwright-crawler/src/internals/adaptive-playwright-crawler.ts b/packages/playwright-crawler/src/internals/adaptive-playwright-crawler.ts
@@ -470,13 +470,16 @@ export class AdaptivePlaywrightCrawler<
         const result = new RequestHandlerResult(this.config, AdaptivePlaywrightCrawler.CRAWLEE_STATE_KEY);
         const logs: LogProxyCall[] = [];
 
+        const deferredCleanup: (() => Promise<unknown>)[] = [];
+
         const resultBoundContextHelpers = {
             addRequests: result.addRequests,
             pushData: result.pushData,
             useState: this.allowStorageAccess(useStateFunction),
             getKeyValueStore: this.allowStorageAccess(result.getKeyValueStore),
             enqueueLinks: result.enqueueLinks,
             log: this.createLogProxy(context.log, logs),
+            registerDeferredCleanup: (cleanup: () => Promise<unknown>) => deferredCleanup.push(cleanup),
         };
 
         try {
@@ -510,6 +513,8 @@ export class AdaptivePlaywrightCrawler<
             return { result, ok: true, logs };
         } catch (error) {
             return { error, ok: false, logs };
+        } finally {
+            await Promise.all(deferredCleanup.map((cleanup) => cleanup()));
         }
     }
 
diff --git a/test/core/crawlers/basic_browser_crawler.ts b/test/core/crawlers/basic_browser_crawler.ts
@@ -1,18 +1,30 @@
 import type { PuppeteerController, PuppeteerPlugin } from '@crawlee/browser-pool';
-import type { PuppeteerCrawlerOptions, PuppeteerCrawlingContext, PuppeteerGoToOptions } from '@crawlee/puppeteer';
+import type {
+    BrowserCrawlerOptions,
+    BrowserCrawlingContext,
+    PuppeteerCrawlingContext,
+    PuppeteerGoToOptions,
+} from '@crawlee/puppeteer';
 import { BrowserCrawler } from '@crawlee/puppeteer';
 import type { HTTPResponse, LaunchOptions, Page } from 'puppeteer';
 
+export type TestCrawlingContext = BrowserCrawlingContext<Page, HTTPResponse, PuppeteerController>;
+
 export class BrowserCrawlerTest extends BrowserCrawler<
     Page,
     HTTPResponse,
     PuppeteerController,
     { browserPlugins: [PuppeteerPlugin] },
     LaunchOptions,
-    PuppeteerCrawlingContext
+    TestCrawlingContext
 > {
-    constructor(options: Partial<PuppeteerCrawlerOptions> = {}) {
-        super(options as any);
+    constructor(
+        options: Partial<BrowserCrawlerOptions<Page, HTTPResponse, PuppeteerController, TestCrawlingContext>> = {},
+    ) {
+        super({
+            ...options,
+            contextPipelineBuilder: () => this.buildContextPipeline(),
+        });
     }
 
     protected async _navigationHandler(
diff --git a/test/core/crawlers/browser_crawler.test.ts b/test/core/crawlers/browser_crawler.test.ts
@@ -2,7 +2,7 @@ import type { Server } from 'node:http';
 
 import { BROWSER_POOL_EVENTS, OperatingSystemsName, PuppeteerPlugin } from '@crawlee/browser-pool';
 import { BLOCKED_STATUS_CODES } from '@crawlee/core';
-import type { PuppeteerCrawlingContext, PuppeteerGoToOptions } from '@crawlee/puppeteer';
+import type { PuppeteerGoToOptions } from '@crawlee/puppeteer';
 import { EnqueueStrategy, ProxyConfiguration, Request, RequestList, RequestState, Session } from '@crawlee/puppeteer';
 import { sleep } from '@crawlee/utils';
 import type { HTTPResponse } from 'puppeteer';
@@ -13,6 +13,7 @@ import { MemoryStorageEmulator } from 'test/shared/MemoryStorageEmulator.js';
 import { ENV_VARS } from '@apify/consts';
 import log from '@apify/log';
 
+import type { TestCrawlingContext } from './basic_browser_crawler.js';
 import { BrowserCrawlerTest } from './basic_browser_crawler.js';
 
 describe('BrowserCrawler', () => {
@@ -64,7 +65,7 @@ describe('BrowserCrawler', () => {
         const processed: Request[] = [];
         const failed: Request[] = [];
         const requestList = await RequestList.open(null, sources);
-        const requestHandler = async ({ page, request, response }: PuppeteerCrawlingContext) => {
+        const requestHandler = async ({ page, request, response }: TestCrawlingContext) => {
             await page.waitForSelector('title');
 
             expect(response!.status()).toBe(200);
@@ -124,7 +125,7 @@ describe('BrowserCrawler', () => {
         let sessionGoto!: Session;
         const browserCrawler = new (class extends BrowserCrawlerTest {
             protected override async _navigationHandler(
-                ctx: PuppeteerCrawlingContext,
+                ctx: TestCrawlingContext,
             ): Promise<HTTPResponse | null | undefined> {
                 vitest.spyOn(ctx.session!, 'markBad');
                 sessionGoto = ctx.session!;
@@ -184,6 +185,7 @@ describe('BrowserCrawler', () => {
             },
             requestList,
             useSessionPool: true,
+            requestHandler: async () => {},
             maxRequestRetries: 0,
             postNavigationHooks: [hook],
         });
@@ -270,7 +272,7 @@ describe('BrowserCrawler', () => {
         let optionsGoto: PuppeteerGoToOptions;
         const browserCrawler = new (class extends BrowserCrawlerTest {
             protected override async _navigationHandler(
-                ctx: PuppeteerCrawlingContext,
+                ctx: TestCrawlingContext,
                 gotoOptions: PuppeteerGoToOptions,
             ): Promise<HTTPResponse | null | undefined> {
                 optionsGoto = gotoOptions;
@@ -617,7 +619,7 @@ describe('BrowserCrawler', () => {
         let called = false;
         const browserCrawler = new (class extends BrowserCrawlerTest {
             protected override async _navigationHandler(
-                ctx: PuppeteerCrawlingContext,
+                ctx: TestCrawlingContext,
             ): Promise<HTTPResponse | null | undefined> {
                 browserCrawler.browserPool.on(BROWSER_POOL_EVENTS.BROWSER_RETIRED, () => {
                     resolve();
@@ -841,7 +843,7 @@ describe('BrowserCrawler', () => {
 
             const browserCrawler = new (class extends BrowserCrawlerTest {
                 protected override async _navigationHandler(
-                    ctx: PuppeteerCrawlingContext,
+                    ctx: TestCrawlingContext,
                 ): Promise<HTTPResponse | null | undefined> {
                     const { session } = ctx;
                     const proxyInfo = await this.proxyConfiguration!.newProxyInfo(session?.id);
@@ -880,7 +882,7 @@ describe('BrowserCrawler', () => {
             let numberOfRotations = -requestList!.length();
             const browserCrawler = new (class extends BrowserCrawlerTest {
                 protected override async _navigationHandler(
-                    ctx: PuppeteerCrawlingContext,
+                    ctx: TestCrawlingContext,
                 ): Promise<HTTPResponse | null | undefined> {
                     const { session } = ctx;
                     const proxyInfo = await this.proxyConfiguration!.newProxyInfo(session?.id);
@@ -918,7 +920,7 @@ describe('BrowserCrawler', () => {
 
             const crawler = new (class extends BrowserCrawlerTest {
                 protected override async _navigationHandler(
-                    ctx: PuppeteerCrawlingContext,
+                    ctx: TestCrawlingContext,
                 ): Promise<HTTPResponse | null | undefined> {
                     const { session } = ctx;
                     const proxyInfo = await this.proxyConfiguration!.newProxyInfo(session?.id);
@@ -967,16 +969,16 @@ describe('BrowserCrawler', () => {
         });
 
         test('uses correct crawling context', async () => {
-            let prepareCrawlingContext: PuppeteerCrawlingContext;
+            let prepareCrawlingContext: TestCrawlingContext;
 
-            const gotoFunction = async (crawlingContext: PuppeteerCrawlingContext) => {
+            const gotoFunction = async (crawlingContext: TestCrawlingContext) => {
                 prepareCrawlingContext = crawlingContext;
                 expect(crawlingContext.request).toBeInstanceOf(Request);
                 expect(crawlingContext.session).toBeInstanceOf(Session);
                 expect(typeof crawlingContext.page).toBe('object');
             };
 
-            const requestHandler = async (crawlingContext: PuppeteerCrawlingContext) => {
+            const requestHandler = async (crawlingContext: TestCrawlingContext) => {
                 expect(crawlingContext === prepareCrawlingContext).toEqual(true);
                 expect(crawlingContext.request).toBeInstanceOf(Request);
                 expect(crawlingContext.session).toBeInstanceOf(Session);
@@ -986,7 +988,7 @@ describe('BrowserCrawler', () => {
                 throw new Error('some error');
             };
 
-            const failedRequestHandler = async (crawlingContext: Partial<PuppeteerCrawlingContext>, error: Error) => {
+            const failedRequestHandler = async (crawlingContext: Partial<TestCrawlingContext>, error: Error) => {
                 expect(crawlingContext).toBe(prepareCrawlingContext);
                 expect(crawlingContext.request).toBeInstanceOf(Request);
                 expect(crawlingContext.session).toBeInstanceOf(Session);
