From 364ff518b7b909761fe5da01db7c0a7242bd70ab Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Sun, 19 Feb 2023 21:40:21 -0500
Subject: [PATCH] Only remove promise in query hook if the subscription was
 removed

---
 .../toolkit/src/query/react/buildHooks.ts     |  4 +-
 .../src/query/tests/buildHooks.test.tsx       | 99 ++++++++++++++++++-
 2 files changed, 99 insertions(+), 4 deletions(-)

diff --git a/packages/toolkit/src/query/react/buildHooks.ts b/packages/toolkit/src/query/react/buildHooks.ts
index 11f48f2c44..9d89188514 100644
--- a/packages/toolkit/src/query/react/buildHooks.ts
+++ b/packages/toolkit/src/query/react/buildHooks.ts
@@ -740,7 +740,9 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       })
 
       usePossiblyImmediateEffect((): void | undefined => {
-        promiseRef.current = undefined
+        if (subscriptionRemoved) {
+          promiseRef.current = undefined
+        }
       }, [subscriptionRemoved])
 
       usePossiblyImmediateEffect((): void | undefined => {
diff --git a/packages/toolkit/src/query/tests/buildHooks.test.tsx b/packages/toolkit/src/query/tests/buildHooks.test.tsx
index feabd12ee5..2ddc91537b 100644
--- a/packages/toolkit/src/query/tests/buildHooks.test.tsx
+++ b/packages/toolkit/src/query/tests/buildHooks.test.tsx
@@ -1,4 +1,3 @@
-import util from 'util'
 import * as React from 'react'
 import type {
   UseMutation,
@@ -10,7 +9,14 @@ import {
   QueryStatus,
   skipToken,
 } from '@reduxjs/toolkit/query/react'
-import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+  renderHook,
+} from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
 import { rest } from 'msw'
 import {
@@ -28,7 +34,6 @@ import type { AnyAction } from 'redux'
 import type { SubscriptionOptions } from '@reduxjs/toolkit/dist/query/core/apiState'
 import type { SerializedError } from '@reduxjs/toolkit'
 import { createListenerMiddleware, configureStore } from '@reduxjs/toolkit'
-import { renderHook } from '@testing-library/react'
 import { delay } from '../../utils'
 
 // Just setup a temporary in-memory counter for tests that `getIncrementedAmount`.
@@ -715,6 +720,94 @@ describe('hooks tests', () => {
       expect(res.data!.amount).toBeGreaterThan(originalAmount)
     })
 
+    // See https://github.com/reduxjs/redux-toolkit/issues/3182
+    test('Hook subscriptions are properly cleaned up when changing skip back and forth', async () => {
+      const pokemonApi = createApi({
+        baseQuery: fetchBaseQuery({ baseUrl: 'https://pokeapi.co/api/v2/' }),
+        endpoints: (builder) => ({
+          getPokemonByName: builder.query({
+            queryFn: (name: string) => ({ data: null }),
+            keepUnusedDataFor: 1,
+          }),
+        }),
+      })
+
+      const storeRef = setupApiStore(pokemonApi, undefined, {
+        withoutTestLifecycles: true,
+      })
+
+      const getSubscriptions = () => storeRef.store.getState().api.subscriptions
+
+      const checkNumSubscriptions = (arg: string, count: number) => {
+        const subscriptions = getSubscriptions()
+        const cacheKeyEntry = subscriptions[arg]
+
+        if (cacheKeyEntry) {
+          expect(Object.values(cacheKeyEntry).length).toBe(count)
+        }
+      }
+
+      // 1) Initial state: an active subscription
+      const { result, rerender, unmount } = renderHook(
+        ([arg, options]: Parameters<
+          typeof pokemonApi.useGetPokemonByNameQuery
+        >) => pokemonApi.useGetPokemonByNameQuery(arg, options),
+        {
+          wrapper: storeRef.wrapper,
+          initialProps: ['a'],
+        }
+      )
+
+      await act(async () => {
+        await delay(1)
+      })
+
+      // 2) Set the current subscription to `{skip: true}
+      await act(async () => {
+        rerender(['a', { skip: true }])
+      })
+
+      // 3) Change _both_ the cache key _and_ `{skip: false}` at the same time.
+      // This causes the `subscriptionRemoved` check to be `true`.
+      await act(async () => {
+        rerender(['b'])
+      })
+
+      // There should only be one active subscription after changing the arg
+      checkNumSubscriptions('b', 1)
+
+      // 4) Re-render with the same arg.
+      // This causes the `subscriptionRemoved` check to be `false`.
+      // Correct behavior is this does _not_ clear the promise ref,
+      // so
+      await act(async () => {
+        rerender(['b'])
+      })
+
+      // There should only be one active subscription after changing the arg
+      checkNumSubscriptions('b', 1)
+
+      await act(async () => {
+        await delay(1)
+      })
+
+      unmount()
+
+      await act(async () => {
+        await delay(1)
+      })
+
+      // There should be no subscription entries left over after changing
+      // cache key args and swapping `skip` on and off
+      checkNumSubscriptions('b', 0)
+
+      const finalSubscriptions = getSubscriptions()
+
+      for (let cacheKeyEntry of Object.values(finalSubscriptions)) {
+        expect(Object.values(cacheKeyEntry!).length).toBe(0)
+      }
+    })
+
     describe('Hook middleware requirements', () => {
       let mock: jest.SpyInstance