From 5f7068ee2287b5a04eda231b4c151b30c17de8d8 Mon Sep 17 00:00:00 2001
From: dvoytenko <dima.voytenko@vercel.com>
Date: Wed, 6 Sep 2023 19:57:48 -0700
Subject: [PATCH] Set test headers via page.router API

---
 .../experimental/testmode/playwright/index.ts |  7 +-----
 .../testmode/playwright/next-fixture.ts       | 15 ++++++-------
 .../testmode/playwright/page-route.ts         | 22 +++++++++++++++----
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/packages/next/src/experimental/testmode/playwright/index.ts b/packages/next/src/experimental/testmode/playwright/index.ts
index 545e7bc303de4..c5b608c48e408 100644
--- a/packages/next/src/experimental/testmode/playwright/index.ts
+++ b/packages/next/src/experimental/testmode/playwright/index.ts
@@ -40,16 +40,11 @@ export const test = base.test.extend<
   ],
 
   next: [
-    async (
-      { nextOptions, _nextWorker, page, extraHTTPHeaders },
-      use,
-      testInfo
-    ) => {
+    async ({ nextOptions, _nextWorker, page }, use, testInfo) => {
       await applyNextFixture(use, {
         testInfo,
         nextWorker: _nextWorker,
         page,
-        extraHTTPHeaders,
         nextOptions,
       })
     },
diff --git a/packages/next/src/experimental/testmode/playwright/next-fixture.ts b/packages/next/src/experimental/testmode/playwright/next-fixture.ts
index aa57a44e6a281..ca68f233ed646 100644
--- a/packages/next/src/experimental/testmode/playwright/next-fixture.ts
+++ b/packages/next/src/experimental/testmode/playwright/next-fixture.ts
@@ -17,9 +17,15 @@ class NextFixtureImpl implements NextFixture {
     private worker: NextWorkerFixture,
     private page: Page
   ) {
+    const testHeaders = {
+      'Next-Test-Proxy-Port': String(worker.proxyPort),
+      'Next-Test-Data': testId,
+    }
     const handleFetch = this.handleFetch.bind(this)
     worker.onFetch(testId, handleFetch)
-    this.page.route('**', (route) => handleRoute(route, page, handleFetch))
+    this.page.route('**', (route) =>
+      handleRoute(route, page, testHeaders, handleFetch)
+    )
   }
 
   teardown(): void {
@@ -51,13 +57,11 @@ export async function applyNextFixture(
     nextOptions,
     nextWorker,
     page,
-    extraHTTPHeaders,
   }: {
     testInfo: TestInfo
     nextOptions: NextOptions
     nextWorker: NextWorkerFixture
     page: Page
-    extraHTTPHeaders: Record<string, string> | undefined
   }
 ): Promise<void> {
   const fixture = new NextFixtureImpl(
@@ -66,11 +70,6 @@ export async function applyNextFixture(
     nextWorker,
     page
   )
-  page.setExtraHTTPHeaders({
-    ...extraHTTPHeaders,
-    'Next-Test-Proxy-Port': String(nextWorker.proxyPort),
-    'Next-Test-Data': fixture.testId,
-  })
 
   await use(fixture)
 
diff --git a/packages/next/src/experimental/testmode/playwright/page-route.ts b/packages/next/src/experimental/testmode/playwright/page-route.ts
index c0b3bd09e3162..1bd083d03fdbe 100644
--- a/packages/next/src/experimental/testmode/playwright/page-route.ts
+++ b/packages/next/src/experimental/testmode/playwright/page-route.ts
@@ -1,16 +1,30 @@
-import type { Page, Route } from '@playwright/test'
+import type { Page, Route, Request } from '@playwright/test'
 import type { FetchHandler } from './next-worker-fixture'
 
+function continueRoute(
+  route: Route,
+  request: Request,
+  testHeaders: Record<string, string>
+): Promise<void> {
+  return route.continue({
+    headers: {
+      ...request.headers(),
+      ...testHeaders,
+    },
+  })
+}
+
 export async function handleRoute(
   route: Route,
   page: Page,
+  testHeaders: Record<string, string>,
   fetchHandler: FetchHandler | null
 ) {
   const request = route.request()
 
   // Continue the navigation and non-fetch requests.
   if (request.isNavigationRequest() || request.resourceType() !== 'fetch') {
-    return route.continue()
+    return continueRoute(route, request, testHeaders)
   }
 
   // Continue the local requests. The followup requests will be intercepted
@@ -18,7 +32,7 @@ export async function handleRoute(
   const pageOrigin = new URL(page.url()).origin
   const requestOrigin = new URL(request.url()).origin
   if (pageOrigin === requestOrigin) {
-    return route.continue()
+    return continueRoute(route, request, testHeaders)
   }
 
   if (!fetchHandler) {
@@ -44,7 +58,7 @@ export async function handleRoute(
     return route.abort()
   }
   if (proxyResponse === 'continue') {
-    return route.continue()
+    return continueRoute(route, request, testHeaders)
   }
   const { status, headers, body } = proxyResponse
   return route.fulfill({