insights/products/wizard/frontend/wizardSessionStreamLogic.test.ts

import { installMockEventSource, MockEventSource } from 'lib/wizard-sync/eventSource.mock'

import { expectLogic } from 'kea-test-utils'
import insights from 'insights-js'

import { ApiError } from 'lib/api-error'
import { FEATURE_FLAGS } from 'lib/constants'
import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
import { SSE_RECONNECT_MAX_MS } from 'lib/wizard-sync/pollLoop'
import { projectLogic } from 'scenes/projectLogic'

import { initKeaTests } from '~/test/init'

import { wizardSessionsLatestRetrieve } from './generated/api'
import type { WizardSessionDTOApi } from './generated/api.schemas'
import { wizardSessionStreamLogic } from './wizardSessionStreamLogic'

jest.mock('./generated/api', () => ({
    wizardSessionsLatestRetrieve: jest.fn(),
    getWizardSessionsStreamRetrieveUrl: jest.fn(() => '/mock-stream-url'),
}))

const mockLatestRetrieve = wizardSessionsLatestRetrieve as jest.Mock

function makeSession(overrides: Partial<WizardSessionDTOApi> = {}): WizardSessionDTOApi {
    return {
        session_id: 'sess-1',
        team_id: 997,
        workflow_id: 'insights-integration',
        skill_id: 'install',
        started_at: '2026-01-01T00:00:00Z',
        run_phase: 'running',
        tasks: [],
        event_plan: null,
        error: null,
        pending_input: null,
        handoff_text: null,
        created_by: null,
        created_at: '2026-01-01T00:00:00Z',
        updated_at: '2026-01-01T00:00:00Z',
        is_stale: false,
        ...overrides,
    }
}

// Max jittered gap for the default 3s interval is 3.6s — advancing past it guarantees the next tick.
const PAST_MAX_JITTERED_INTERVAL_MS = 4000

describe('wizardSessionStreamLogic polling mode', () => {
    let logic: ReturnType<typeof wizardSessionStreamLogic.build>

    beforeEach(async () => {
        initKeaTests()
        mockLatestRetrieve.mockReset()
        logic = wizardSessionStreamLogic({ workflowId: 'insights-integration' })
        logic.mount()
        featureFlagLogic.actions.setFeatureFlags([FEATURE_FLAGS.ONBOARDING_WIZARD_SYNC_MODE], {
            [FEATURE_FLAGS.ONBOARDING_WIZARD_SYNC_MODE]: 'polling',
        })
        // connect() no-ops into an error without a project — wait for the test bootstrap to provide one.
        await expectLogic(projectLogic).toMatchValues({ currentProjectId: expect.any(Number) })
        jest.useFakeTimers()
    })

    afterEach(() => {
        logic?.unmount()
        jest.useRealTimers()
    })

    it('polls the latest session and feeds it through sessionUpdated', async () => {
        const session = makeSession()
        mockLatestRetrieve.mockResolvedValue(session)

        logic.actions.connect()
        await expectLogic(logic)
            .toDispatchActions(['connectionOpened', 'sessionUpdated'])
            .toMatchValues({ latestSession: session, connectionStatus: 'open' })

        expect(mockLatestRetrieve).toHaveBeenCalledWith(
            expect.any(String),
            {
                workflow_id: 'insights-integration',
                skill_id: undefined,
            },
            { headers: { 'X-Wizard-Poll-Source': 'transport' } }
        )
    })

    it('keeps polling through a 204 (no session yet) without clobbering latestSession', async () => {
        // The generated client resolves an empty 204 body to null (getJSONOrNull), not undefined.
        mockLatestRetrieve.mockResolvedValue(null)

        logic.actions.connect()
        await expectLogic(logic)
            .toDispatchActions(['connectionOpened'])
            .toNotHaveDispatchedActions(['sessionUpdated'])
            .toMatchValues({ latestSession: null, connectionStatus: 'open' })

        jest.advanceTimersByTime(PAST_MAX_JITTERED_INTERVAL_MS)
        await Promise.resolve()
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(2)
    })

    it('stops polling on disconnect', async () => {
        mockLatestRetrieve.mockResolvedValue(makeSession())

        logic.actions.connect()
        await expectLogic(logic).toDispatchActions(['sessionUpdated'])
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)

        logic.actions.disconnect()
        jest.advanceTimersByTime(PAST_MAX_JITTERED_INTERVAL_MS)
        await Promise.resolve()
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)
    })

    it('stops polling permanently on a 404', async () => {
        mockLatestRetrieve.mockRejectedValue(new ApiError('not found', 404))

        logic.actions.connect()
        await expectLogic(logic).toDispatchActions(['connectionErrored']).toMatchValues({ connectionStatus: 'error' })
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)

        jest.advanceTimersByTime(10 * PAST_MAX_JITTERED_INTERVAL_MS)
        await Promise.resolve()
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)
    })

    it('backs off after a transient error instead of retrying at full cadence', async () => {
        mockLatestRetrieve.mockRejectedValueOnce(new Error('network hiccup')).mockResolvedValue(makeSession())

        logic.actions.connect()
        await expectLogic(logic).toDispatchActions(['connectionErrored'])
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)

        // One failure doubles the 3s base to 6s (±20% jitter → at least 4.8s): the base window
        // must NOT produce a retry...
        jest.advanceTimersByTime(PAST_MAX_JITTERED_INTERVAL_MS)
        await Promise.resolve()
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(1)

        // ...but the backed-off window must.
        jest.advanceTimersByTime(PAST_MAX_JITTERED_INTERVAL_MS)
        await expectLogic(logic).toDispatchActions(['sessionUpdated'])
        expect(mockLatestRetrieve).toHaveBeenCalledTimes(2)
    })
})

// The first backoff step is 2s ±20%, so this window always contains the retry.
const PAST_FIRST_RECONNECT_MS = 2600

describe('wizardSessionStreamLogic SSE mode', () => {
    let logic: ReturnType<typeof wizardSessionStreamLogic.build>
    let restoreEventSource: () => void
    let capture: jest.SpyInstance

    beforeEach(async () => {
        // Resolved feature flags are persisted, so the polling suite above would otherwise decide
        // this suite's transport too.
        window.localStorage.clear()
        initKeaTests()
        mockLatestRetrieve.mockReset()
        restoreEventSource = installMockEventSource()
        capture = jest.spyOn(insights, 'capture').mockImplementation(() => undefined as never)
        logic = wizardSessionStreamLogic({ workflowId: 'insights-integration' })
        logic.mount()
        // connect() no-ops into an error without a project — wait for the test bootstrap to provide one.
        await expectLogic(projectLogic).toMatchValues({ currentProjectId: expect.any(Number) })
        jest.useFakeTimers()
    })

    afterEach(() => {
        logic?.unmount()
        jest.useRealTimers()
        restoreEventSource()
        capture.mockRestore()
    })

    const transportErrors = (): unknown[] =>
        capture.mock.calls.filter(([event]) => event === 'wizard sync transport error')

    it('reconnects after a fatal close and reports the transport error once per connect cycle', async () => {
        logic.actions.connect()
        MockEventSource.last().emitOpen()
        await expectLogic(logic).toDispatchActions(['connectionOpened'])

        MockEventSource.last().emitTransportError(MockEventSource.CLOSED)
        await expectLogic(logic).toDispatchActions(['connectionErrored'])
        jest.advanceTimersByTime(PAST_FIRST_RECONNECT_MS)
        expect(MockEventSource.instances).toHaveLength(2)

        // A stream that stays fatally closed retries for as long as the tab is open. One capture per
        // backoff step would be one every 30s per tab, enough to drown the transport comparison.
        MockEventSource.last().emitTransportError(MockEventSource.CLOSED)
        await expectLogic(logic).toDispatchActions(['connectionErrored'])
        expect(transportErrors()).toHaveLength(1)
    })

    it('cancels a pending reconnect on disconnect', async () => {
        logic.actions.connect()
        MockEventSource.last().emitTransportError(MockEventSource.CLOSED)
        await expectLogic(logic).toDispatchActions(['connectionErrored'])

        logic.actions.disconnect()
        jest.advanceTimersByTime(SSE_RECONNECT_MAX_MS)
        expect(MockEventSource.instances).toHaveLength(1)
    })
})