fix: reset generator state on .resetHandlers()/.restoreHandlers()

Author: kettanaito

src/core/experimental/define-network.ts

@@ -47,8 +47,9 @@ export interface DefineNetworkOptions<
   onUnhandledFrame?: UnhandledFrameHandle
 }
 
-export interface NetworkApi<Sources extends Array<NetworkSource<any>>>
-  extends NetworkHandlersApi {
+export interface NetworkApi<
+  Sources extends Array<NetworkSource<any>>,
+> extends NetworkHandlersApi {
   readyState: NetworkReadyState
   /**
    * Enable the network interception and handling.
@@ -210,11 +211,7 @@ export function defineNetwork<Sources extends Array<NetworkSource<any>>>(
       handlersController.reset(handlers)
     },
     restoreHandlers() {
-      for (const handler of handlersController.currentHandlers()) {
-        if ('isUsed' in handler) {
-          handler.isUsed = false
-        }
-      }
+      handlersController.restore()
     },
     listHandlers() {
       return toReadonlyArray(handlersController.currentHandlers())

src/core/experimental/handlers-controller.ts

@@ -90,6 +90,12 @@ export abstract class HandlersController {
       ),
     )
 
+    for (const handler of this.currentHandlers()) {
+      if ('reset' in handler) {
+        handler['reset']()
+      }
+    }
+
     const { initialHandlers } = this.getState()
 
     if (nextHandlers.length === 0) {
@@ -108,6 +114,14 @@ export abstract class HandlersController {
     })
   }
 
+  public restore(): void {
+    for (const handler of this.currentHandlers()) {
+      if ('restore' in handler) {
+        handler['restore']()
+      }
+    }
+  }
+
   #validateHandlers(handlers: Array<AnyHandler>): boolean {
     return handlers.every((handler) => !Array.isArray(handler))
   }

src/core/experimental/setup-api.ts

@@ -42,11 +42,7 @@ export abstract class SetupApi<
   }
 
   public restoreHandlers(): void {
-    this.handlersController.currentHandlers().forEach((handler) => {
-      if ('isUsed' in handler) {
-        handler.isUsed = false
-      }
-    })
+    this.handlersController.restore()
   }
 
   public resetHandlers(...nextHandlers: Array<AnyHandler>): void {

src/core/handlers/RequestHandler.ts

@@ -171,6 +171,27 @@ export abstract class RequestHandler<
     this.isUsed = false
   }
 
+  /**
+   * Reset the runtime state accumulated during response resolution,
+   * such as generator iterator progress. Called when this handler is
+   * removed from the active handlers list so re-adding it later starts
+   * from a clean state.
+   */
+  protected reset(): void {
+    this.resolverIterator = undefined
+    this.resolverIteratorResult = undefined
+  }
+
+  /**
+   * Restore this handler so it can match requests again after being
+   * exhausted (e.g. via `{ once: true }`). Also clears any accumulated
+   * resolution state.
+   */
+  protected restore(): void {
+    this.reset()
+    this.isUsed = false
+  }
+
   /**
    * Determine if the intercepted request should be mocked.
    */

src/core/utils/internal/requestHandlerUtils.ts

@@ -1,6 +1,4 @@
-/**
- * @vitest-environment node
- */
+// @vitest-environment node
 import { http, HttpResponse, delay } from 'msw'
 import { setupServer } from 'msw/node'
 
@@ -24,7 +22,7 @@ afterAll(() => {
 
 it('supports generator function as response resolver', async () => {
   server.use(
-    http.get('https://example.com/weather', function* () {
+    http.get('http://localhost/weather', function* () {
       let degree = 10
 
       while (degree < 13) {
@@ -38,18 +36,18 @@ it('supports generator function as response resolver', async () => {
   )
 
   // Must respond with yielded responses.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(11)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(12)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(13)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(11)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(12)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(13)
   // Must respond with the final "done" response.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(14)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(14)
   // Must keep responding with the final "done" response.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(14)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(14)
 })
 
 it('supports async generator function as response resolver', async () => {
   server.use(
-    http.get('https://example.com/weather', async function* () {
+    http.get('http://localhost/weather', async function* () {
       await delay(20)
 
       let degree = 10
@@ -64,17 +62,17 @@ it('supports async generator function as response resolver', async () => {
     }),
   )
 
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(11)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(12)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(13)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(14)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(14)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(11)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(12)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(13)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(14)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(14)
 })
 
 it('supports generator function as one-time response resolver', async () => {
   server.use(
     http.get(
-      'https://example.com/weather',
+      'http://localhost/weather',
       function* () {
         let degree = 10
 
@@ -94,16 +92,31 @@ it('supports generator function as one-time response resolver', async () => {
   )
 
   // Must respond with the yielded incrementing responses.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(11)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(12)
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(13)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(11)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(12)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(13)
   // Must respond with the "done" final response from the iterator.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(14)
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe(14)
   // Must respond with the other handler since the generator one is used.
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(
-    'fallback',
-  )
-  await expect(fetchJson('https://example.com/weather')).resolves.toEqual(
-    'fallback',
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe('fallback')
+  await expect(fetchJson('http://localhost/weather')).resolves.toBe('fallback')
+})
+
+it('resets the generator state after the handlers are reset', async () => {
+  server.use(
+    http.get('http://localhost/resource', function* () {
+      yield HttpResponse.json('Yield')
+      return HttpResponse.json('Stable')
+    }),
   )
+
+  await server.boundary(async () => {
+    await expect(fetchJson('http://localhost/resource')).resolves.toBe('Yield')
+    await expect(fetchJson('http://localhost/resource')).resolves.toBe('Stable')
+
+    server.resetHandlers()
+
+    await expect(fetchJson('http://localhost/resource')).resolves.toBe('Yield')
+    await expect(fetchJson('http://localhost/resource')).resolves.toBe('Stable')
+  })()
 })