CHIM-1187 Preserve school profile selection flow

src/components/ProfileMenu/ProfileMenu.test.tsx

@@ -7,11 +7,14 @@ import { Util } from '../../utility/util';
 import { useAppSelector } from '../../redux/hooks';
 import { MemoryRouter } from 'react-router-dom';
 import {
+  CURRENT_MODE,
   EVENTS,
+  MODES,
   PAGES,
   STICKER_BOOK_NOTIFICATION_DOT_ENABLED,
   ENABLE_STICKER_BOOK,
 } from '../../common/constants';
+import { schoolUtil } from '../../utility/schoolUtil';
 
 // --- Mocks ---
 
@@ -36,15 +39,22 @@ jest.mock('../../i18n', () => ({
 }));
 
 const mockPush = jest.fn();
+const mockReplace = jest.fn();
 jest.mock('react-router', () => ({
   ...jest.requireActual('react-router'),
   useHistory: () => ({
     push: mockPush,
-    replace: jest.fn(),
+    replace: mockReplace,
     location: { pathname: '/' },
   }),
 }));
 
+jest.mock('../../utility/schoolUtil', () => ({
+  schoolUtil: {
+    setCurrentClass: jest.fn(),
+  },
+}));
+
 // Mock useGbContext
 jest.mock('../../growthbook/Growthbook', () => ({
   useGbContext: () => ({ setGbUpdated: jest.fn() }),
@@ -60,10 +70,12 @@ const mockApi = {
 };
 
 const mockStudent = { id: 'student-123', name: 'Test Student' };
+const mockSetCurrentClass = schoolUtil.setCurrentClass as jest.Mock;
 
 describe('ProfileMenu Notification Logic', () => {
   beforeEach(() => {
     jest.clearAllMocks();
+    localStorage.clear();
     mockApi.getUserStickerBook.mockReset();
     mockApi.getUserStickerBook.mockResolvedValue([]);
     mockApi.markStciekercolledasTrue.mockReset();
@@ -191,4 +203,30 @@ describe('ProfileMenu Notification Logic', () => {
       expect.any(Object),
     );
   });
+
+  test('keeps school-mode switch profile on SelectMode selection flow', async () => {
+    localStorage.setItem(CURRENT_MODE, MODES.TEACHER_SCHOOL);
+
+    render(
+      <MemoryRouter>
+        <ProfileMenu onClose={jest.fn()} />
+      </MemoryRouter>,
+    );
+
+    await waitFor(() => expect(mockApi.getUserStickerBook).toHaveBeenCalled());
+
+    const switchProfileItem = screen
+      .getByText(/Switch Profile/i)
+      .closest('.profile-menu-item');
+    fireEvent.click(switchProfileItem!);
+
+    await waitFor(() => {
+      expect(Util.setCurrentStudent).toHaveBeenCalledWith(null);
+    });
+    expect(mockSetCurrentClass).not.toHaveBeenCalled();
+    expect(mockReplace).toHaveBeenCalledWith(PAGES.SELECT_MODE, {
+      from: '/',
+      fromSchoolModeSwitchProfile: true,
+    });
+  });
 });

src/components/ProfileMenu/ProfileMenu.tsx

@@ -177,6 +177,7 @@ const ProfileMenu = ({ onClose }: ProfileMenuProps) => {
     setGbUpdated(true);
     history.replace(PAGES.SELECT_MODE, {
       from: history.location.pathname,
+      fromSchoolModeSwitchProfile: true,
     });
   };
 

src/pages/SelectMode.test.tsx

@@ -70,7 +70,12 @@ jest.mock('./assets/leftArrowIcon.svg', () => ({
 }));
 
 const mockHistoryReplace = jest.fn();
-let mockLocationState: { fromKidsAppLocationSchool?: boolean } | undefined;
+let mockLocationState:
+  | {
+      fromKidsAppLocationSchool?: boolean;
+      fromSchoolModeSwitchProfile?: boolean;
+    }
+  | undefined;
 jest.mock('react-router', () => {
   const actual = jest.requireActual('react-router');
   return {
@@ -914,6 +919,41 @@ describe('SelectMode page', () => {
     expect(mockHistoryReplace).not.toHaveBeenCalledWith(PAGES.DISPLAY_SCHOOLS);
   });
 
+  it('keeps school-mode switch profile in school flow for teacher-role users', async () => {
+    const teacherSchool = { id: 'school-1', name: 'Teacher School' };
+    const classDoc = {
+      id: 'class-1',
+      name: 'Class 1',
+      school_id: teacherSchool.id,
+    };
+    mockLocationState = { fromSchoolModeSwitchProfile: true };
+    mockGetCurrMode.mockResolvedValue(MODES.TEACHER_SCHOOL);
+    mockAuthHandler.getCurrentUser.mockResolvedValue({
+      id: 'user-1',
+      name: 'Teacher User',
+    });
+    mockApiHandler.getSchoolsForUser.mockResolvedValue([
+      { school: teacherSchool, role: 'TEACHER' },
+    ]);
+    mockApiHandler.getSchoolsWithRoleAutouser.mockResolvedValue([]);
+    mockApiHandler.getClassesForSchool.mockResolvedValue([classDoc]);
+    mockApiHandler.getStudentsForClass.mockResolvedValue([
+      { id: 'student-1', name: 'Student 1' },
+    ]);
+
+    render(<SelectMode />);
+
+    await waitFor(() => {
+      expect(mockApiHandler.getClassesForSchool).toHaveBeenCalledWith(
+        teacherSchool.id,
+        'user-1',
+      );
+    });
+    expect(mockSetCurrMode).not.toHaveBeenCalledWith(MODES.TEACHER);
+    expect(mockHistoryReplace).not.toHaveBeenCalledWith(PAGES.HOME_PAGE);
+    expect(mockHistoryReplace).not.toHaveBeenCalledWith(PAGES.DISPLAY_SCHOOLS);
+  });
+
   it('uses the school picker when multiple teacher schools remain after stored school is removed', async () => {
     const removedPrincipalSchool = {
       id: 'school-removed',

src/pages/SelectMode.tsx

@@ -101,6 +101,7 @@ interface SchoolModeOption {
 
 interface SelectModeLocationState {
   fromKidsAppLocationSchool?: boolean;
+  fromSchoolModeSwitchProfile?: boolean;
 }
 
 const SUPPORTED_LANGUAGE_CODES = new Set<string>(Object.values(LANG));
@@ -450,7 +451,8 @@ const SelectMode: FC = () => {
 
     const shouldSuppressTeacherAutoEntry =
       currentMode === MODES.TEACHER_SCHOOL &&
-      location.state?.fromKidsAppLocationSchool === true;
+      (location.state?.fromKidsAppLocationSchool === true ||
+        location.state?.fromSchoolModeSwitchProfile === true);
     const shouldAutoEnterTeacherApp =
       teacherRoleEntries.length > 0 && !shouldSuppressTeacherAutoEntry;
     const shouldUseEmptySchoolFallback = !shouldSuppressTeacherAutoEntry;