fix: ensure optimistic cache updates for unenroll and saved for later (#1331)

adamstankiewicz · web-flow · commit 397fcbb5004e · 2025-03-10T13:23:00.000-04:00
diff --git a/src/components/app/routes/data/utils.js b/src/components/app/routes/data/utils.js
@@ -87,13 +87,13 @@ export async function ensureEnterpriseAppData({
           }
 
           // Optimistically update the query cache with the auto-activated or auto-applied subscription license.
-          queryClient.setQueryData(subscriptionsQuery.queryKey, {
-            ...queryClient.getQueryData(subscriptionsQuery.queryKey),
+          queryClient.setQueryData(subscriptionsQuery.queryKey, (oldData) => ({
+            ...oldData,
             subscriptionLicensesByStatus: updatedLicensesByStatus,
             subscriptionPlan: activatedOrAutoAppliedLicense.subscriptionPlan,
             subscriptionLicense: activatedOrAutoAppliedLicense,
             subscriptionLicenses: updatedSubscriptionLicenses,
-          });
+          }));
         }
 
         return subscriptionsData;
@@ -331,14 +331,14 @@ export async function ensureActiveEnterpriseCustomerUser({
     enterpriseCustomer = nextActiveEnterpriseCustomer;
     activeEnterpriseCustomer = nextActiveEnterpriseCustomer;
     allLinkedEnterpriseCustomerUsers = updatedLinkedEnterpriseCustomerUsers;
-    // Optimistically update the BFF layer (use helper)
+    // Optimistically update the BFF query cache
     if (matchedBFFQuery) {
-      queryClient.setQueryData(matchedBFFQuery({ enterpriseSlug }), {
-        ...queryClient.getQueryData(matchedBFFQuery({ enterpriseSlug })),
+      queryClient.setQueryData(matchedBFFQuery({ enterpriseSlug }).queryKey, (oldData) => ({
+        ...oldData,
         enterpriseCustomer,
         activeEnterpriseCustomer,
         allLinkedEnterpriseCustomerUsers: updatedLinkedEnterpriseCustomerUsers,
-      });
+      }));
     }
     return {
       enterpriseCustomer,
diff --git a/src/components/app/routes/data/utils.test.js b/src/components/app/routes/data/utils.test.js
@@ -194,22 +194,19 @@ describe('ensureActiveEnterpriseCustomerUser', () => {
           });
           if (isBFFData) {
             expect(mockQueryClient.setQueryData).toHaveBeenCalledTimes(timesUpdateActiveEnterpriseCustomerCalled);
-            expect(mockQueryClient.setQueryData).toHaveBeenCalledWith(
-              {
-                queryFn: expect.any(Function),
-                queryKey: queryEnterpriseLearnerDashboardBFF(
-                  { enterpriseSlug: expectedEnterpriseCustomer.slug },
-                ).queryKey,
-              },
-              {
-                ...mockQueryClient.getQueryData(queryEnterpriseLearnerDashboardBFF(
-                  { enterpriseSlug: expectedEnterpriseCustomer.slug },
-                )),
-                enterpriseCustomer: expectedEnterpriseCustomer,
-                activeEnterpriseCustomer: expectedEnterpriseCustomer,
-                allLinkedEnterpriseCustomerUsers: expectedAllLinkedEnterpriseCustomers,
-              },
-            );
+            const expectedQuery = queryEnterpriseLearnerDashboardBFF({
+              enterpriseSlug: expectedEnterpriseCustomer.slug,
+            });
+            expect(mockQueryClient.setQueryData).toHaveBeenCalledWith(expectedQuery.queryKey, expect.any(Function));
+            // Get the function that was passed to setQueryData
+            const updateFunction = mockQueryClient.setQueryData.mock.calls[0][1];
+            const result = updateFunction(mockEnterpriseLearnerData);
+            expect(result).toEqual({
+              enterpriseCustomer: expectedEnterpriseCustomer,
+              activeEnterpriseCustomer: expectedEnterpriseCustomer,
+              allLinkedEnterpriseCustomerUsers: expectedAllLinkedEnterpriseCustomers,
+              shouldUpdateActiveEnterpriseCustomerUser: false,
+            });
           }
         }
       },
diff --git a/src/components/app/routes/loaders/tests/rootLoader.test.jsx b/src/components/app/routes/loaders/tests/rootLoader.test.jsx
@@ -431,7 +431,19 @@ describe('rootLoader', () => {
         });
 
         // Assert the subscriptions query cache is optimistically updated
-        expect(mockQueryClient.setQueryData).toHaveBeenCalledWith(subscriptionsQuery.queryKey, {
+        expect(mockQueryClient.setQueryData).toHaveBeenCalledWith(subscriptionsQuery.queryKey, expect.any(Function));
+        // Get the function that was passed to setQueryData
+        const updateFunction = mockQueryClient.setQueryData.mock.calls[0][1];
+        // Call the function with a mock oldData value if needed to simulate the cache update
+        const result = updateFunction({
+          subscriptionLicenses: [mockSubscriptionsData.subscriptionLicense],
+          subscriptionLicensesByStatus: {
+            [LICENSE_STATUS.PENDING]: [mockSubscriptionsData.subscriptionLicense],
+          },
+          subscriptionLicense: mockSubscriptionsData.subscriptionLicense,
+          subscriptionPlan: mockSubscriptionsData.subscriptionPlan,
+        });
+        expect(result).toEqual({
           subscriptionLicenses: [
             {
               ...mockSubscriptionsData.subscriptionLicense,
diff --git a/src/components/dashboard/main-content/course-enrollments/course-cards/unenroll/UnenrollModal.jsx b/src/components/dashboard/main-content/course-enrollments/course-cards/unenroll/UnenrollModal.jsx
@@ -4,7 +4,7 @@ import { useQueryClient } from '@tanstack/react-query';
 import {
   ActionRow, Alert, AlertModal, Button, StatefulButton,
 } from '@openedx/paragon';
-import { logError, logInfo } from '@edx/frontend-platform/logging';
+import { logError } from '@edx/frontend-platform/logging';
 
 import { ToastsContext } from '../../../../../Toasts';
 import { unenrollFromCourse } from './data';
@@ -51,28 +51,29 @@ const UnenrollModal = ({
       const bffQueryKeysToUpdate = [dashboardBFFQueryKey];
       // Update the enterpriseCourseEnrollments data in the cache for each BFF query.
       bffQueryKeysToUpdate.forEach((queryKey) => {
-        const existingBFFData = queryClient.getQueryData(queryKey);
-        if (!existingBFFData) {
-          logInfo(`Skipping optimistic cache update of ${JSON.stringify(queryKey)} as no cached query data exists yet.`);
-          return;
-        }
-        const updatedBFFData = {
-          ...existingBFFData,
-          enterpriseCourseEnrollments: existingBFFData.enterpriseCourseEnrollments.filter(enrollmentForCourseFilter),
-        };
-        queryClient.setQueryData(queryKey, updatedBFFData);
+        queryClient.setQueryData(queryKey, (oldData) => {
+          if (!oldData) {
+            return oldData;
+          }
+          return {
+            ...oldData,
+            enterpriseCourseEnrollments: oldData.enterpriseCourseEnrollments.filter(enrollmentForCourseFilter),
+            allEnrollmentsByStatus: Object.keys(oldData.allEnrollmentsByStatus).reduce((acc, status) => {
+              const filteredEnrollments = oldData.allEnrollmentsByStatus[status].filter(enrollmentForCourseFilter);
+              acc[status] = filteredEnrollments;
+              return acc;
+            }, {}),
+          };
+        });
       });
     }
 
     // Update the legacy queryEnterpriseCourseEnrollments cache as well.
     const enterpriseCourseEnrollmentsQueryKey = queryEnterpriseCourseEnrollments(enterpriseCustomer.uuid).queryKey;
-    const existingCourseEnrollmentsData = queryClient.getQueryData(enterpriseCourseEnrollmentsQueryKey);
-    if (!existingCourseEnrollmentsData) {
-      logInfo(`Skipping optimistic cache update of ${JSON.stringify(enterpriseCourseEnrollmentsQueryKey)} as no cached query data exists yet.`);
-      return;
-    }
-    const updatedCourseEnrollmentsData = existingCourseEnrollmentsData.filter(enrollmentForCourseFilter);
-    queryClient.setQueryData(enterpriseCourseEnrollmentsQueryKey, updatedCourseEnrollmentsData);
+    queryClient.setQueryData(enterpriseCourseEnrollmentsQueryKey, (oldData) => {
+      const updatedCourseEnrollmentsData = oldData?.filter(enrollmentForCourseFilter);
+      return updatedCourseEnrollmentsData;
+    });
   };
 
   const handleUnenrollButtonClick = async () => {
diff --git a/src/components/dashboard/main-content/course-enrollments/course-cards/unenroll/UnenrollModal.test.jsx b/src/components/dashboard/main-content/course-enrollments/course-cards/unenroll/UnenrollModal.test.jsx
@@ -2,7 +2,6 @@ import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { QueryClientProvider } from '@tanstack/react-query';
 import '@testing-library/jest-dom/extend-expect';
-import { logInfo } from '@edx/frontend-platform/logging';
 import { COURSE_STATUSES } from '../../../../../../constants';
 import { unenrollFromCourse } from './data';
 import UnenrollModal from './UnenrollModal';
@@ -35,10 +34,6 @@ jest.mock('../../../../../app/data', () => ({
   fetchEnterpriseLearnerDashboard: jest.fn(),
 }));
 
-jest.mock('@edx/frontend-platform/logging', () => ({
-  logInfo: jest.fn(),
-}));
-
 const mockEnterpriseCustomer = enterpriseCustomerFactory();
 const mockEnterpriseCourseEnrollment = enterpriseCourseEnrollmentFactory();
 const mockEnterpriseCourseEnrollments = [mockEnterpriseCourseEnrollment];
@@ -188,14 +183,11 @@ describe('<UnenrollModal />', () => {
       const bffDashboardData = mockQueryClient.getQueryData(
         queryEnterpriseLearnerDashboardBFF({ enterpriseSlug: mockEnterpriseCustomer.slug }).queryKey,
       );
-      let expectedLogInfoCalls = 0;
       if (isBFFEnabled) {
         // Only verify the BFF queryEnterpriseCourseEnrollments cache is updated if BFF feature is enabled.
         let expectedBFFDashboardData;
         if (existingBFFDashboardQueryData) {
           expectedBFFDashboardData = learnerDashboardBFFResponse;
-        } else {
-          expectedLogInfoCalls += 1;
         }
         expect(bffDashboardData).toEqual(expectedBFFDashboardData);
       } else {
@@ -214,14 +206,9 @@ describe('<UnenrollModal />', () => {
       let expectedLegacyEnrollmentsData;
       if (existingEnrollmentsQueryData) {
         expectedLegacyEnrollmentsData = [];
-      } else {
-        expectedLogInfoCalls += 1;
       }
       expect(legacyEnrollmentsData).toEqual(expectedLegacyEnrollmentsData);
 
-      // Verify logInfo calls
-      expect(logInfo).toHaveBeenCalledTimes(expectedLogInfoCalls);
-
       // Verify side effects
       expect(mockOnSuccess).toHaveBeenCalledTimes(1);
       expect(mockAddToast).toHaveBeenCalledTimes(1);
diff --git a/src/components/dashboard/main-content/course-enrollments/data/hooks.js b/src/components/dashboard/main-content/course-enrollments/data/hooks.js
@@ -4,7 +4,7 @@ import {
 import { useMutation, useQueryClient } from '@tanstack/react-query';
 import { AppContext } from '@edx/frontend-platform/react';
 import { camelCaseObject } from '@edx/frontend-platform/utils';
-import { logError, logInfo } from '@edx/frontend-platform/logging';
+import { logError } from '@edx/frontend-platform/logging';
 import { sendEnterpriseTrackEventWithDelay } from '@edx/frontend-enterprise-utils';
 
 import { useLocation } from 'react-router-dom';
@@ -42,6 +42,7 @@ import { ASSIGNMENTS_EXPIRING_WARNING_LOCALSTORAGE_KEY } from '../../../data/con
 import { LICENSE_STATUS } from '../../../../enterprise-user-subsidy/data/constants';
 import { useStatefulEnroll } from '../../../../stateful-enroll/data';
 import { COURSE_STATUSES } from '../../../../../constants';
+import { findCourseStatusKey } from '../../../../../utils/common';
 
 /**
  * Return data for upgrading a course using the user's subsidies
@@ -494,32 +495,47 @@ export function useUpdateCourseEnrollmentStatus() {
       const dashboardBFFQueryKey = queryEnterpriseLearnerDashboardBFF({
         enterpriseSlug: enterpriseCustomer.slug,
       }).queryKey;
-
       const bffQueryKeysToUpdate = [dashboardBFFQueryKey];
       // Update the enterpriseCourseEnrollments data in the cache for each BFF query.
       bffQueryKeysToUpdate.forEach((queryKey) => {
-        const existingBFFData = queryClient.getQueryData(queryKey);
-        if (!existingBFFData) {
-          logInfo(`Skipping optimistic cache update of ${JSON.stringify(queryKey)} as no cached query data exists yet.`);
-          return;
-        }
-        const updatedBFFData = {
-          ...existingBFFData,
-          enterpriseCourseEnrollments: existingBFFData.enterpriseCourseEnrollments.map(transformUpdatedEnrollment),
-        };
-        queryClient.setQueryData(queryKey, updatedBFFData);
+        queryClient.setQueryData(queryKey, (oldData) => {
+          if (!oldData) {
+            return oldData;
+          }
+          const updatedEnrollments = oldData.enterpriseCourseEnrollments.map(transformUpdatedEnrollment);
+          const updatedAllEnrollmentsByStatus = Object.keys(oldData.allEnrollmentsByStatus).reduce((acc, status) => {
+            acc[status] = oldData.allEnrollmentsByStatus[status]
+              .filter((enrollment) => enrollment.courseRunId !== courseRunId)
+              .map(transformUpdatedEnrollment);
+            return acc;
+          }, {});
+
+          // Find the updated enrollment.
+          const updatedEnrollment = updatedEnrollments.find(enrollment => enrollment.courseRunId === courseRunId);
+          if (updatedEnrollment) {
+            const newCourseStatusKey = findCourseStatusKey(updatedEnrollment.courseRunStatus);
+            // Add the enrollment to the new status group.
+            if (!updatedAllEnrollmentsByStatus[newCourseStatusKey]) {
+              updatedAllEnrollmentsByStatus[newCourseStatusKey] = [];
+            }
+            updatedAllEnrollmentsByStatus[newCourseStatusKey].push(updatedEnrollment);
+          }
+
+          return {
+            ...oldData,
+            enterpriseCourseEnrollments: updatedEnrollments,
+            allEnrollmentsByStatus: updatedAllEnrollmentsByStatus,
+          };
+        });
       });
     }
 
     // Update the legacy queryEnterpriseCourseEnrollments cache as well.
     const enterpriseCourseEnrollmentsQueryKey = queryEnterpriseCourseEnrollments(enterpriseCustomer.uuid).queryKey;
-    const existingCourseEnrollmentsData = queryClient.getQueryData(enterpriseCourseEnrollmentsQueryKey);
-    if (!existingCourseEnrollmentsData) {
-      logInfo(`Skipping optimistic cache update of ${JSON.stringify(enterpriseCourseEnrollmentsQueryKey)} as no cached query data exists yet.`);
-      return;
-    }
-    const updatedCourseEnrollmentsData = existingCourseEnrollmentsData.map(transformUpdatedEnrollment);
-    queryClient.setQueryData(enterpriseCourseEnrollmentsQueryKey, updatedCourseEnrollmentsData);
+    queryClient.setQueryData(enterpriseCourseEnrollmentsQueryKey, (oldData) => {
+      const updatedData = oldData?.map(transformUpdatedEnrollment);
+      return updatedData;
+    });
   }, [queryClient, enterpriseCustomer, isBFFEnabled]);
 }
 
diff --git a/src/components/dashboard/main-content/course-enrollments/data/tests/hooks.test.jsx b/src/components/dashboard/main-content/course-enrollments/data/tests/hooks.test.jsx
@@ -1,6 +1,5 @@
 import { act, renderHook } from '@testing-library/react-hooks';
 import * as logger from '@edx/frontend-platform/logging';
-import { logInfo } from '@edx/frontend-platform/logging';
 import { AppContext } from '@edx/frontend-platform/react';
 import { sendEnterpriseTrackEventWithDelay } from '@edx/frontend-enterprise-utils';
 import dayjs from 'dayjs';
@@ -1218,7 +1217,6 @@ describe('useUpdateCourseEnrollmentStatus', () => {
           enterpriseSlug: mockEnterpriseCustomer.slug,
         }).queryKey,
       );
-      let expectedLogInfoCalls = 0;
       const expectedCourseRunStatus = doesCourseRunIdMatch
         ? newEnrollmentStatus
         : originalEnrollmentStatus;
@@ -1229,9 +1227,6 @@ describe('useUpdateCourseEnrollmentStatus', () => {
         );
         if (existingBFFDashboardQueryData) {
           expect(foundMockEnrollment.courseRunStatus).toEqual(expectedCourseRunStatus);
-        } else {
-          expectedLogInfoCalls += 1;
-          expect(dashboardBFFData).toBeUndefined();
         }
       }
 
@@ -1245,12 +1240,8 @@ describe('useUpdateCourseEnrollmentStatus', () => {
       if (existingEnrollmentsQueryData) {
         expect(foundMockEnrollment.courseRunStatus).toEqual(expectedCourseRunStatus);
       } else {
-        expectedLogInfoCalls += 1;
         expect(enrollmentsData).toBeUndefined();
       }
-
-      // Verify logInfo calls
-      expect(logInfo).toHaveBeenCalledTimes(expectedLogInfoCalls);
     });
   });
 });
diff --git a/src/utils/common.js b/src/utils/common.js
@@ -1,5 +1,6 @@
 import { logError } from '@edx/frontend-platform/logging';
 import dayjs from './dayjs';
+import { COURSE_STATUSES } from '../constants';
 
 export const isCourseEnded = endDate => dayjs(endDate).isBefore(dayjs());
 
@@ -193,3 +194,12 @@ export const formatPrice = (price, options = {}) => {
   });
   return USDollar.format(Math.abs(price));
 };
+
+export function findCourseStatusKey(statusValue) {
+  for (const key in COURSE_STATUSES) {
+    if (COURSE_STATUSES[key] === statusValue) {
+      return key;
+    }
+  }
+  return undefined;
+}
