Testmode: intercept rewrite fetches (#54259)

The test mode should allow interception of all fetches made by the
server, including rewrites.

---------

Co-authored-by: JJ Kasper <jj@jjsweb.site>

packages/next/src/experimental/testmode/playwright/msw.ts

@@ -61,12 +61,17 @@ export const test = base.extend<{
           referrer,
           referrerPolicy,
         })
+        let isUnhandled = false
         let isPassthrough = false
         let mockedResponse: MockedResponse | undefined
         await handleRequest(
           mockedRequest,
           handlers.slice(0),
-          { onUnhandledRequest: 'error' },
+          {
+            onUnhandledRequest: () => {
+              isUnhandled = true
+            },
+          },
           emitter as any,
           {
             onPassthroughResponse: () => {
@@ -78,6 +83,9 @@ export const test = base.extend<{
           }
         )
 
+        if (isUnhandled) {
+          return undefined
+        }
         if (isPassthrough) {
           return 'continue'
         }

packages/next/src/experimental/testmode/server.ts

@@ -1,11 +1,12 @@
 import { AsyncLocalStorage } from 'async_hooks'
-import { NodeRequestHandler } from '../../server/next-server'
 import type {
   ProxyFetchRequest,
   ProxyFetchResponse,
   ProxyResponse,
 } from './proxy'
 import { ClientRequestInterceptor } from 'next/dist/compiled/@mswjs/interceptors/ClientRequest'
+import { WorkerRequestHandler } from '../../server/lib/setup-server-worker'
+import { NodeRequestHandler } from '../../server/next-server'
 
 interface TestReqInfo {
   url: string
@@ -92,7 +93,9 @@ async function handleFetch(
       return originalFetch(request)
     case 'abort':
     case 'unhandled':
-      throw new Error('Proxy request aborted')
+      throw new Error(
+        `Proxy request aborted [${request.method} ${request.url}]`
+      )
     default:
       break
   }
@@ -132,7 +135,29 @@ export function interceptTestApis(): () => void {
   }
 }
 
-export function wrapRequestHandler(
+export function wrapRequestHandlerWorker(
+  handler: WorkerRequestHandler
+): WorkerRequestHandler {
+  return async (req, res) => {
+    const proxyPortHeader = req.headers['next-test-proxy-port']
+    if (!proxyPortHeader) {
+      await handler(req, res)
+      return
+    }
+
+    const url = req.url ?? ''
+    const proxyPort = Number(proxyPortHeader)
+    const testData = (req.headers['next-test-data'] as string | undefined) ?? ''
+    const testReqInfo: TestReqInfo = {
+      url,
+      proxyPort,
+      testData,
+    }
+    await testStorage.run(testReqInfo, () => handler(req, res))
+  }
+}
+
+export function wrapRequestHandlerNode(
   handler: NodeRequestHandler
 ): NodeRequestHandler {
   return async (req, res, parsedUrl) => {

packages/next/src/server/lib/router-server.ts

@@ -302,7 +302,7 @@ export async function initialize(opts: {
     devInstance?.ensureMiddleware
   )
 
-  const requestHandler: WorkerRequestHandler = async (req, res) => {
+  const requestHandlerImpl: WorkerRequestHandler = async (req, res) => {
     if (compress) {
       // @ts-expect-error not express req/res
       compress(req, res, () => {})
@@ -718,6 +718,17 @@ export async function initialize(opts: {
     }
   }
 
+  let requestHandler: WorkerRequestHandler = requestHandlerImpl
+  if (opts.experimentalTestProxy) {
+    // Intercept fetch and other testmode apis.
+    const {
+      wrapRequestHandlerWorker,
+      interceptTestApis,
+    } = require('../../experimental/testmode/server')
+    requestHandler = wrapRequestHandlerWorker(requestHandler)
+    interceptTestApis()
+  }
+
   const upgradeHandler: WorkerUpgradeHandler = async (req, socket, head) => {
     try {
       req.on('error', (_err) => {

packages/next/src/server/next-server.ts

@@ -1073,8 +1073,10 @@ export default class NextNodeServer extends BaseServer {
   public getRequestHandler(): NodeRequestHandler {
     const handler = this.makeRequestHandler()
     if (this.serverOptions.experimentalTestProxy) {
-      const { wrapRequestHandler } = require('../experimental/testmode/server')
-      return wrapRequestHandler(handler)
+      const {
+        wrapRequestHandlerNode,
+      } = require('../experimental/testmode/server')
+      return wrapRequestHandlerNode(handler)
     }
     return handler
   }