FI-943 fix: map headers for redirected requests also

src/context/waitForEventsState.ts

@@ -5,6 +5,7 @@ import type {
   AllRequestsCompletePredicateWithPromise,
   RequestPredicateWithPromise,
   ResponsePredicateWithPromise,
+  Url,
   WaitForEventsState,
 } from '../types/internal';
 import type {RequestHookToWaitForEvents} from '../utils/requestHooks';
@@ -34,6 +35,7 @@ export const getWaitForEventsState = (
       null,
     ) as WaitForEventsState['hashOfNotCompleteRequests'],
     hook: {} as RequestHookToWaitForEvents,
+    redirects: Object.create(null) as Record<Url, Url>,
     requestPredicates: new Set<RequestPredicateWithPromise>(),
     responsePredicates: new Set<ResponsePredicateWithPromise>(),
   };

src/types/waitForEvents.ts

@@ -2,7 +2,7 @@ import type {RequestHookToWaitForEvents} from '../utils/requestHooks';
 
 import type {UtcTimeInMs} from './date';
 import type {MergeFunctions} from './fn';
-import type {Request, RequestWithUtcTimeInMs, Response, ResponseWithRequest} from './http';
+import type {Request, RequestWithUtcTimeInMs, Response, ResponseWithRequest, Url} from './http';
 import type {MaybePromise} from './promise';
 import type {RequestHookContextId} from './requestHooks';
 
@@ -92,6 +92,7 @@ export type WaitForEventsState = Readonly<{
   allRequestsCompletePredicates: Set<AllRequestsCompletePredicateWithPromise>;
   hashOfNotCompleteRequests: Record<RequestHookContextId, RequestWithUtcTimeInMs>;
   hook: RequestHookToWaitForEvents;
+  redirects: Record<Url, Url>;
   requestPredicates: Set<RequestPredicateWithPromise>;
   responsePredicates: Set<ResponsePredicateWithPromise>;
 }>;

src/utils/clientFunction/clientFunctionWrapper.ts

@@ -7,7 +7,7 @@ declare const originalFn: (...args: unknown[]) => OriginalClientFunctionResult;
 declare const printedClientFunctionName: string;
 
 /**
- * This client function wraps all ClientFunction bodies and maps them errors to error messages.
+ * This client function wraps all `ClientFunction` bodies and maps them errors to error messages.
  * @internal
  */
 export const clientFunctionWrapper = <Args extends readonly unknown[], R>(

src/utils/requestHooks/RequestHookToWaitForEvents.ts

@@ -12,6 +12,7 @@ import {getDurationWithUnits} from '../getDurationWithUnits';
 import {mapBackendResponseForLogs} from '../log';
 import {addNotCompleteRequest, completeRequest, processEventsPredicates} from '../waitForEvents';
 
+import {addRedirectToWaitForEventsState} from './addRedirectToWaitForEventsState';
 import {getHeaderValue} from './getHeaderValue';
 import {getRequestFromRequestOptions} from './getRequestFromRequestOptions';
 import {getResponseFromResponseEvent} from './getResponseFromResponseEvent';
@@ -47,6 +48,7 @@ export class RequestHookToWaitForEvents extends RequestHookWithEvents {
 
     cdpClient.on('Network.requestWillBeSent', ({redirectResponse}) => {
       if (redirectResponse) {
+        addRedirectToWaitForEventsState(redirectResponse, waitForEventsState);
         removeNotCompleteRequestsByUrl(redirectResponse.url as Url, waitForEventsState);
       }
     });

src/utils/requestHooks/SetHeadersRequestHook.ts

@@ -6,6 +6,7 @@ import {
   REQUEST_HOOK_CONTEXT_KEY,
   RESOLVED_PROMISE,
 } from '../../constants/internal';
+import {getWaitForEventsState} from '../../context/waitForEventsState';
 import {testController} from '../../testController';
 
 import {assertValueIsBoolean, assertValueIsDefined} from '../asserts';
@@ -15,6 +16,7 @@ import {setReadonlyProperty} from '../setReadonlyProperty';
 import {applyHeadersMapper} from './applyHeadersMapper';
 import {applyHeadersMapperOnCdpMode} from './applyHeadersMapperOnCdpMode';
 import {getHeadersFromHeaderEntries} from './getHeadersFromHeaderEntries';
+import {RequestHookToWaitForEvents} from './RequestHookToWaitForEvents';
 import {RequestHookWithEvents} from './RequestHookWithEvents';
 
 import type {
@@ -34,7 +36,24 @@ export class SetHeadersRequestHook extends RequestHookWithEvents {
     private readonly url: Url,
     private readonly options: MapOptions,
   ) {
-    super([url], INCLUDE_HEADERS_IN_RESPONSE_EVENT);
+    const waitForEventsState = getWaitForEventsState(RequestHookToWaitForEvents);
+    let wasCalled = false;
+
+    const predicate = (request: Readonly<{url?: string}>): boolean => {
+      if (request.url === url) {
+        wasCalled = true;
+
+        return true;
+      }
+
+      if (wasCalled && url in waitForEventsState.redirects) {
+        return waitForEventsState.redirects[url] === request.url;
+      }
+
+      return false;
+    };
+
+    super(predicate, INCLUDE_HEADERS_IN_RESPONSE_EVENT);
   }
 
   override onRequest(event: RequestHookRequestEvent): Promise<void> {

src/utils/requestHooks/addRedirectToWaitForEventsState.ts

@@ -0,0 +1,19 @@
+import type Protocol from 'devtools-protocol';
+
+import type {Url, WaitForEventsState} from '../../types/internal';
+
+/**
+ * Adds (register) new redirect into `WaitForEventsState`, if any.
+ * @internal
+ */
+export const addRedirectToWaitForEventsState = (
+  redirectResponse: Protocol.Network.Response,
+  waitForEventsState: WaitForEventsState,
+): void => {
+  const {location} = redirectResponse.headers;
+
+  if (location !== undefined) {
+    // eslint-disable-next-line no-param-reassign
+    waitForEventsState.redirects[redirectResponse.url as Url] = location as Url;
+  }
+};