Websockets Implementation for the Connecting screen

README.md

@@ -48,6 +48,8 @@ const App = () => {
 | `profiles`                     | [`ProfilesTypes`](./typings/connectProps.d.ts)             | The connect widget uses the profiles to set the initial state of the widget. [More details](./docs/PROFILES.md)                                                                  | See more details                              |
 | `userFeatures`                 | [`UserFeaturesType`](./typings/connectProps.d.ts)          | The connect widget uses user features to determine the behavior of the widget. [More details](./docs/USER_FEATURES.md)                                                           | See more details                              |
 | `showTooSmallDialog`           | `boolean`                                                  | The connect widget can show a warning when the widget size is below the supported 320px.                                                                                         | `true`                                        |
+| `webSocketConnection`          | `object`                                                   | An object containing `isConnected()` function and `webSocketMessages$` observable for real-time updates.                                                                         | `null`                                        |
+| `experimentalFeatures`         | `object`                                                   | An object to enable or disable experimental features like `useWebSockets: true`.                                                                                                 | `null`                                        |
 
 ## ApiProvider
 

src/ConnectWidget.tsx

@@ -9,6 +9,7 @@ import { initGettextLocaleData } from 'src/utilities/Personalization'
 import { ConnectedTokenProvider } from 'src/ConnectedTokenProvider'
 import { TooSmallDialog } from 'src/components/app/TooSmallDialog'
 import { setLocalizedContent } from 'src/redux/reducers/localizedContentSlice'
+import { WebSocketProvider } from 'src/context/WebSocketContext'
 import './sharedVariables.css'
 
 interface PostMessageContextType {
@@ -27,6 +28,7 @@ export const ConnectWidget = ({
   onAnalyticPageview = () => {},
   postMessageEventOverrides,
   showTooSmallDialog = true,
+  webSocketConnection,
   ...props
 }: any) => {
   initGettextLocaleData(props.language)
@@ -38,12 +40,14 @@ export const ConnectWidget = ({
   return (
     <Provider store={Store}>
       <ConnectedTokenProvider>
-        <PostMessageContext.Provider value={{ onPostMessage, postMessageEventOverrides }}>
-          <WidgetDimensionObserver heightOffset={0}>
-            {showTooSmallDialog && <TooSmallDialog onAnalyticPageview={onAnalyticPageview} />}
-            <Connect onAnalyticPageview={onAnalyticPageview} {...props} />
-          </WidgetDimensionObserver>
-        </PostMessageContext.Provider>
+        <WebSocketProvider value={webSocketConnection}>
+          <PostMessageContext.Provider value={{ onPostMessage, postMessageEventOverrides }}>
+            <WidgetDimensionObserver heightOffset={0}>
+              {showTooSmallDialog && <TooSmallDialog onAnalyticPageview={onAnalyticPageview} />}
+              <Connect onAnalyticPageview={onAnalyticPageview} {...props} />
+            </WidgetDimensionObserver>
+          </PostMessageContext.Provider>
+        </WebSocketProvider>
       </ConnectedTokenProvider>
     </Provider>
   )

src/__tests__/ConnectWidget-test.tsx

@@ -0,0 +1,53 @@
+import React from 'react'
+import { render } from '@testing-library/react'
+import { describe, it, expect, vi } from 'vitest'
+import { of } from 'rxjs'
+
+import { ConnectWidget } from '../ConnectWidget'
+import { useWebSocket } from '../context/WebSocketContext'
+
+vi.mock('src/Connect', () => ({
+  default: vi.fn(() => {
+    // In actual implementation, it uses Context
+    // But for the test we just want to see if it renders without crashing
+    // and correctly provides the context which we can check via useWebSocket in a child if we want
+    return <div data-test="mock-connect">mock-connect</div>
+  }),
+}))
+
+// A simple component to verify context
+const ContextChecker = () => {
+  const ws = useWebSocket()
+  return <div data-test="context-checker">{ws ? 'has-ws' : 'no-ws'}</div>
+}
+
+// We need to mock Connect to render the ContextChecker instead
+vi.mock('src/Connect', () => ({
+  default: () => <ContextChecker />,
+}))
+
+describe('ConnectWidget', () => {
+  const defaultProps = {
+    clientConfig: {},
+    profiles: {},
+    userFeatures: {},
+    language: { locale: 'en', localizedContent: {} },
+  }
+
+  it('provides webSocketConnection to children when passed as a prop', () => {
+    const mockWS = {
+      isConnected: vi.fn().mockReturnValue(true),
+      webSocketMessages$: of({}),
+    }
+
+    const { getByTestId } = render(<ConnectWidget {...defaultProps} webSocketConnection={mockWS} />)
+
+    expect(getByTestId('context-checker')).toHaveTextContent('has-ws')
+  })
+
+  it('does not provide webSocketConnection when not passed', () => {
+    const { getByTestId } = render(<ConnectWidget {...defaultProps} />)
+
+    expect(getByTestId('context-checker')).toHaveTextContent('no-ws')
+  })
+})

src/context/WebSocketContext.tsx

@@ -0,0 +1,19 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import React, { createContext, useContext } from 'react'
+import { Observable } from 'rxjs'
+
+export interface WebSocketConnection {
+  isConnected: () => boolean
+  webSocketMessages$: Observable<any>
+}
+
+const WebSocketContext = createContext<WebSocketConnection | undefined>(undefined)
+
+export const WebSocketProvider: React.FC<{
+  value?: WebSocketConnection
+  children: React.ReactNode
+}> = ({ value, children }) => (
+  <WebSocketContext.Provider value={value}>{children}</WebSocketContext.Provider>
+)
+
+export const useWebSocket = () => useContext(WebSocketContext)

src/hooks/__tests__/usePollMember-test.tsx

@@ -1,21 +1,29 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
 import React from 'react'
 import { renderHook, waitFor } from '@testing-library/react'
 import { vi } from 'vitest'
 import { usePollMember, PollingState } from 'src/hooks/usePollMember'
 import { ApiProvider, ApiContextTypes } from 'src/context/ApiContext'
+import { WebSocketProvider, WebSocketConnection } from 'src/context/WebSocketContext'
 import { Provider } from 'react-redux'
 import { createReduxStore, RootState } from 'src/redux/Store'
 import { member, JOB_DATA } from 'src/services/mockedData'
 import { ReadableStatuses } from 'src/const/Statuses'
 import { CONNECTING_MESSAGES } from 'src/utilities/pollers'
 import { take } from 'rxjs/operators'
+import { Subject } from 'rxjs'
 
-const createWrapper = (apiValue: Partial<ApiContextTypes>, preloadedState?: Partial<RootState>) => {
+const createWrapper = (
+  apiValue: Partial<ApiContextTypes>,
+  preloadedState?: Partial<RootState>,
+  webSocketValue?: WebSocketConnection,
+) => {
   const store = createReduxStore(preloadedState)
   const Wrapper = ({ children }: { children: React.ReactNode }) => (
     <Provider store={store}>
-      {/* eslint-disable-next-line @typescript-eslint/no-explicit-any */}
-      <ApiProvider apiValue={apiValue as any}>{children}</ApiProvider>
+      <WebSocketProvider value={webSocketValue}>
+        <ApiProvider apiValue={apiValue as any}>{children}</ApiProvider>
+      </WebSocketProvider>
     </Provider>
   )
   Wrapper.displayName = 'TestWrapper'
@@ -27,6 +35,10 @@ describe('usePollMember', () => {
     document.documentElement.setAttribute('lang', 'en')
   })
 
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
   it('should return a pollMember function', () => {
     const apiValue = {
       loadMemberByGuid: vi.fn().mockResolvedValue(member.member),
@@ -303,9 +315,12 @@ describe('usePollMember', () => {
   })
 
   it('should increment pollingCount on each poll', async () => {
+    const member1 = { ...member.member, guid: 'MBR-1', most_recent_job_guid: 'JOB-1' }
+    const member2 = { ...member.member, guid: 'MBR-2', most_recent_job_guid: 'JOB-2' }
+
     const apiValue = {
-      loadMemberByGuid: vi.fn().mockResolvedValue(member.member),
-      loadJob: vi.fn().mockResolvedValue(JOB_DATA),
+      loadMemberByGuid: vi.fn().mockResolvedValueOnce(member1).mockResolvedValue(member2),
+      loadJob: vi.fn().mockImplementation((guid) => Promise.resolve({ ...JOB_DATA, guid })),
     }
 
     const preloadedState = {
@@ -446,15 +461,23 @@ describe('usePollMember', () => {
       async_account_data_ready: true,
     }
 
-    const memberWithJob = {
+    const member1 = {
       ...member.member,
+      guid: 'MBR-1',
+      most_recent_job_guid: 'JOB-1',
       is_being_aggregated: false,
       connection_status: ReadableStatuses.CONNECTED,
     }
+    const member2 = { ...member1, guid: 'MBR-2', most_recent_job_guid: 'JOB-2' }
+    const member3 = { ...member1, guid: 'MBR-3', most_recent_job_guid: 'JOB-3' }
 
     const apiValue = {
-      loadMemberByGuid: vi.fn().mockResolvedValue(memberWithJob),
-      loadJob: vi.fn().mockResolvedValue(jobWithAsyncData),
+      loadMemberByGuid: vi
+        .fn()
+        .mockResolvedValueOnce(member1)
+        .mockResolvedValueOnce(member2)
+        .mockResolvedValue(member3),
+      loadJob: vi.fn().mockImplementation((guid) => Promise.resolve({ ...jobWithAsyncData, guid })),
     }
 
     const preloadedState = {
@@ -622,12 +645,12 @@ describe('usePollMember', () => {
   }, 10000)
 
   it('should correctly update previousResponse and currentResponse over multiple polls', async () => {
-    const member1 = { ...member.member, guid: 'MBR-1' }
-    const member2 = { ...member.member, guid: 'MBR-2' }
+    const member1 = { ...member.member, guid: 'MBR-1', most_recent_job_guid: 'JOB-1' }
+    const member2 = { ...member.member, guid: 'MBR-2', most_recent_job_guid: 'JOB-2' }
 
     const apiValue = {
       loadMemberByGuid: vi.fn().mockResolvedValueOnce(member1).mockResolvedValue(member2),
-      loadJob: vi.fn().mockResolvedValue(JOB_DATA),
+      loadJob: vi.fn().mockImplementation((guid) => Promise.resolve({ ...JOB_DATA, guid })),
     }
 
     const preloadedState = {
@@ -658,25 +681,34 @@ describe('usePollMember', () => {
 
     // First poll
     expect(states[0].previousResponse).toEqual({})
-    expect(states[0].currentResponse).toEqual({ member: member1, job: JOB_DATA })
+    expect(states[0].currentResponse).toEqual({
+      member: member1,
+      job: { ...JOB_DATA, guid: 'JOB-1' },
+    })
 
     // Second poll
-    expect(states[1].previousResponse).toEqual({ member: member1, job: JOB_DATA })
-    expect(states[1].currentResponse).toEqual({ member: member2, job: JOB_DATA })
+    expect(states[1].previousResponse).toEqual({
+      member: member1,
+      job: { ...JOB_DATA, guid: 'JOB-1' },
+    })
+    expect(states[1].currentResponse).toEqual({
+      member: member2,
+      job: { ...JOB_DATA, guid: 'JOB-2' },
+    })
 
     subscription.unsubscribe()
   }, 10000)
 
   it('should preserve previousResponse and currentResponse when an intermediate poll fails', async () => {
-    const member1 = { ...member.member, guid: 'MBR-1' }
+    const member1 = { ...member.member, guid: 'MBR-1', most_recent_job_guid: 'JOB-1' }
 
     const apiValue = {
       loadMemberByGuid: vi
         .fn()
         .mockResolvedValueOnce(member1)
         .mockRejectedValueOnce(new Error('Intermediate Error'))
-        .mockResolvedValue(member1),
-      loadJob: vi.fn().mockResolvedValue(JOB_DATA),
+        .mockResolvedValue({ ...member1, guid: 'MBR-1-new', most_recent_job_guid: 'JOB-1-new' }),
+      loadJob: vi.fn().mockImplementation((guid) => Promise.resolve({ ...JOB_DATA, guid })),
     }
 
     const preloadedState = {
@@ -707,18 +739,88 @@ describe('usePollMember', () => {
 
     // First poll: Success
     expect(states[0].isError).toBe(false)
-    expect(states[0].currentResponse).toEqual({ member: member1, job: JOB_DATA })
+    expect(states[0].currentResponse).toEqual({
+      member: member1,
+      job: { ...JOB_DATA, guid: 'JOB-1' },
+    })
 
     // Second poll: Error
     expect(states[1].isError).toBe(true)
     expect(states[1].previousResponse).toEqual({}) // Should be preserved from acc
-    expect(states[1].currentResponse).toEqual({ member: member1, job: JOB_DATA }) // Should be preserved from acc
+    expect(states[1].currentResponse).toEqual({
+      member: member1,
+      job: { ...JOB_DATA, guid: 'JOB-1' },
+    }) // Should be preserved from acc
 
     // Third poll: Success again
     expect(states[2].isError).toBe(false)
-    expect(states[2].previousResponse).toEqual({ member: member1, job: JOB_DATA }) // acc.currentResponse was preserved
-    expect(states[2].currentResponse).toEqual({ member: member1, job: JOB_DATA })
+    expect(states[2].previousResponse).toEqual({
+      member: member1,
+      job: { ...JOB_DATA, guid: 'JOB-1' },
+    }) // acc.currentResponse was preserved
+    expect(states[2].currentResponse).toEqual({
+      member: { ...member1, guid: 'MBR-1-new', most_recent_job_guid: 'JOB-1-new' },
+      job: { ...JOB_DATA, guid: 'JOB-1-new' },
+    })
 
     subscription.unsubscribe()
   }, 10000)
+
+  it('should receive updates from WebSockets when enabled', async () => {
+    const wsMessages$ = new Subject<any>()
+    const mockWS = {
+      isConnected: vi.fn().mockReturnValue(true),
+      webSocketMessages$: wsMessages$.asObservable(),
+    }
+
+    const apiValue = {
+      loadMemberByGuid: vi.fn().mockResolvedValue(member.member),
+      loadJob: vi.fn().mockResolvedValue(JOB_DATA),
+    }
+
+    const preloadedState = {
+      experimentalFeatures: {
+        useWebSockets: true,
+        memberPollingMilliseconds: 10000, // Long interval to avoid poll interference
+      },
+    }
+
+    const { result } = renderHook(() => usePollMember(), {
+      wrapper: createWrapper(apiValue, preloadedState, mockWS),
+    })
+
+    const pollMember = result.current
+    const states: PollingState[] = []
+
+    const subscription = pollMember('MBR-123').subscribe((state: PollingState) => {
+      states.push(state)
+    })
+
+    // Emit from WebSocket
+    const wsMember = { guid: 'MBR-123', connection_status: 1 }
+    wsMessages$.next({ event: 'members/updated', payload: wsMember })
+
+    await waitFor(
+      () => {
+        expect(states.length).toBeGreaterThan(0)
+      },
+      { timeout: 4000 },
+    )
+
+    expect(states[0].currentResponse?.member).toEqual(wsMember)
+
+    // Emit priority data ready
+    wsMessages$.next({ event: 'members/priority_data_ready', payload: wsMember })
+
+    await waitFor(
+      () => {
+        expect(states.length).toBeGreaterThan(1)
+      },
+      { timeout: 4000 },
+    )
+
+    expect(states[1].initialDataReady).toBe(true)
+
+    subscription.unsubscribe()
+  })
 })

src/hooks/usePollMember.tsx

@@ -1,6 +1,7 @@
 import { useMemo } from 'react'
 import { DEFAULT_POLLING_STATE, handlePollingResponse } from 'src/utilities/pollers'
 import { useApi } from 'src/context/ApiContext'
+import { useWebSocket } from 'src/context/WebSocketContext'
 import { useSelector } from 'react-redux'
 import { getExperimentalFeatures } from 'src/redux/reducers/experimentalFeaturesSlice'
 
@@ -22,12 +23,13 @@ export interface PollingState {
 
 export function usePollMember() {
   const { api } = useApi()
+  const webSocket = useWebSocket()
 
   const clientLocale = useMemo(() => {
     return document.querySelector('html')?.getAttribute('lang') || 'en'
   }, [document.querySelector('html')?.getAttribute('lang')])
 
-  const { optOutOfEarlyUserRelease, memberPollingMilliseconds } =
+  const { optOutOfEarlyUserRelease, memberPollingMilliseconds, useWebSockets } =
     useSelector(getExperimentalFeatures)
 
   const pollingInterval = memberPollingMilliseconds || 3000
@@ -46,7 +48,9 @@ export function usePollMember() {
       {
         pollingInterval,
         clientLocale,
+        useWebSockets,
       },
+      webSocket,
     )
 
     return updateStream$.pipe(
@@ -72,7 +76,6 @@ export function usePollMember() {
           if (
             !isError &&
             !acc.initialDataReady &&
-            // @ts-expect-error response might be undefined or an error
             response?.job?.async_account_data_ready &&
             !optOutOfEarlyUserRelease
           ) {