mattermost-mobile-upstream/app/managers/websocket_manager.test.ts

// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.txt for license information.

import NetInfo from '@react-native-community/netinfo';
import {AppState} from 'react-native';
import {BackgroundTimer} from 'react-native-nitro-bg-timer-plus';

import {fetchStatusByIds} from '@actions/remote/user';
import {handleFirstConnect, handleReconnect} from '@actions/websocket';
import {hasActiveNativeCall} from '@calls/native_call_mappings';
import WebSocketClient from '@client/websocket';
import DatabaseManager from '@database/manager';
import {launchMark} from '@init/launch_profiler';
import {getCurrentUserId} from '@queries/servers/system';
import {queryAllUsers} from '@queries/servers/user';
import TestHelper from '@test/test_helper';
import {logError} from '@utils/log';

import WebsocketManager from './websocket_manager';

import type {ServerDatabase} from '@typings/database/database';

jest.mock('@react-native-community/netinfo');
jest.mock('react-native/Libraries/AppState/AppState');
jest.mock('@actions/local/user');
jest.mock('@actions/remote/user');
jest.mock('@actions/websocket');
jest.mock('@actions/websocket/event');
jest.mock('@calls/native_call_mappings', () => ({
    hasActiveNativeCall: jest.fn().mockReturnValue(false),
}));
jest.mock('@client/websocket');
jest.mock('@database/manager');
jest.mock('@init/launch_profiler', () => ({
    launchMark: jest.fn(),
}));
jest.mock('@queries/servers/system');
jest.mock('@queries/servers/user');
jest.mock('@utils/log');

describe('WebsocketManager', () => {
    let manager: typeof WebsocketManager;
    let mockWebSocketClient: any;
    let mockCallbacks: {[key: string]: (...args: any[]) => void};
    const mockServerUrl = 'https://example.com';
    const mockToken = 'mock-token';
    const mockCredentials = [{serverUrl: mockServerUrl, token: mockToken} as ServerCredential];

    beforeEach(async () => {
        mockCallbacks = {};
        await DatabaseManager.init([mockServerUrl]);

        // Reset NetInfo mock
        (NetInfo.fetch as jest.Mock).mockResolvedValue({isConnected: true, type: 'wifi'});
        (NetInfo.addEventListener as jest.Mock).mockReturnValue(jest.fn());

        // Reset AppState mock
        AppState.currentState = 'active';

        // Reset WebSocketClient mock
        mockWebSocketClient = {
            setFirstConnectCallback: jest.fn((cb) => {
                mockCallbacks.firstConnect = cb;
            }),

            setEventCallback: jest.fn(),
            setReconnectCallback: jest.fn((cb) => {
                mockCallbacks.reconnect = cb;
            }),
            setReliableReconnectCallback: jest.fn(),
            setCloseCallback: jest.fn((cb) => {
                mockCallbacks.close = cb;
            }),
            initialize: jest.fn(),
            isConnected: jest.fn().mockReturnValue(true),
            close: jest.fn(),
            invalidate: jest.fn(),
            waitForClose: jest.fn().mockResolvedValue(undefined),
        };
        (WebSocketClient as jest.Mock).mockImplementation(() => mockWebSocketClient);

        // Reset DatabaseManager mock
        jest.spyOn(DatabaseManager, 'getServerDatabaseAndOperator').mockImplementation(() => ({
            database: {},
            operator: {},
        } as ServerDatabase));

        manager = WebsocketManager;
    });

    afterEach(async () => {
        await DatabaseManager.destroyServerDatabase(mockServerUrl);
        jest.clearAllMocks();
    });

    describe('init', () => {
        it('should initialize correctly', async () => {
            await manager.init(mockCredentials);

            expect(NetInfo.fetch).toHaveBeenCalled();
            expect(WebSocketClient).toHaveBeenCalledWith(mockServerUrl, mockToken, undefined);
            expect(NetInfo.addEventListener).toHaveBeenCalled();
        });

        it('should handle initialization error gracefully', async () => {
            (DatabaseManager.getServerDatabaseAndOperator as jest.Mock).mockImplementationOnce(() => {
                throw new Error('Database error');
            });

            await manager.init(mockCredentials);

            expect(jest.mocked(logError)).toHaveBeenCalled();
        });
    });

    describe('client management', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
            jest.clearAllMocks();
        });

        it('should create and invalidate clients correctly', async () => {
            const client = await manager.createClient(mockServerUrl, mockToken);
            expect(client).toBeDefined();

            await manager.invalidateClient(mockServerUrl);
            expect(manager.getClient(mockServerUrl)).toBeUndefined();
        });

        it('should handle websocket state observations', () => {
            const observable = manager.observeWebsocketState(mockServerUrl);
            expect(observable).toBeDefined();
        });
    });

    describe('closeAll', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
            jest.clearAllMocks();
        });

        it('should close clients without calling invalidate', () => {
            manager.closeAll();

            expect(mockWebSocketClient.close).toHaveBeenCalledWith(true);
            expect(mockWebSocketClient.invalidate).not.toHaveBeenCalled();
        });
    });

    describe('invalidateClient', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
            jest.clearAllMocks();
        });

        it('should call close, waitForClose, then invalidate in order', async () => {
            const callOrder: string[] = [];
            mockWebSocketClient.close.mockImplementation(() => callOrder.push('close'));
            mockWebSocketClient.waitForClose.mockImplementation(() => {
                callOrder.push('waitForClose');
                return Promise.resolve();
            });
            mockWebSocketClient.invalidate.mockImplementation(() => callOrder.push('invalidate'));

            await manager.invalidateClient(mockServerUrl);

            expect(callOrder).toEqual(['close', 'waitForClose', 'invalidate']);
        });

        it('should handle missing client gracefully', async () => {
            await manager.invalidateClient('https://nonexistent.com');
            expect(mockWebSocketClient.close).not.toHaveBeenCalled();
        });
    });

    describe('connection handling', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
            jest.mocked(handleFirstConnect).mockResolvedValue(undefined);
        });

        it('should handle first connect correctly', async () => {
            mockWebSocketClient.isConnected.mockReturnValueOnce(false);

            await manager.initializeClient(mockServerUrl);

            expect(mockWebSocketClient.initialize).toHaveBeenCalledWith({skipConfigWait: false}, true);
            expect(handleFirstConnect).toHaveBeenCalledWith(mockServerUrl, 'WebSocket Reconnect');

            if (mockCallbacks.firstConnect) {
                mockCallbacks.firstConnect();
            }
        });

        it('should pass skipConfigWait on first connect and mark ws_connecting', async () => {
            mockWebSocketClient.isConnected.mockReturnValue(false);

            await manager.initializeClient(mockServerUrl, 'Cold Start', {skipConfigWait: true});

            expect(mockWebSocketClient.initialize).toHaveBeenCalledWith({skipConfigWait: true}, true);
            expect(handleFirstConnect).toHaveBeenCalledWith(mockServerUrl, 'Cold Start');
            expect(launchMark).toHaveBeenCalledWith('ws_connecting');
        });

        it('should still run handleFirstConnect when already connected but not yet synced', async () => {
            mockWebSocketClient.isConnected.mockReturnValue(true);

            await manager.initializeClient(mockServerUrl, 'Cold Start');

            expect(mockWebSocketClient.initialize).not.toHaveBeenCalled();
            expect(handleFirstConnect).toHaveBeenCalledWith(mockServerUrl, 'Cold Start');
        });

        it('should return the same in-flight promise for concurrent initializeClient calls', async () => {
            mockWebSocketClient.isConnected.mockReturnValue(false);
            let release = () => {};
            jest.mocked(handleFirstConnect).mockImplementation(() => new Promise((resolve) => {
                release = () => resolve(undefined);
            }));

            const first = manager.initializeClient(mockServerUrl, 'Cold Start');
            const second = manager.initializeClient(mockServerUrl, 'Login');

            expect(second).toBe(first);
            expect(handleFirstConnect).toHaveBeenCalledTimes(1);
            expect(handleFirstConnect).toHaveBeenCalledWith(mockServerUrl, 'Cold Start');
            expect(mockWebSocketClient.initialize).toHaveBeenCalledTimes(1);

            release();
            await Promise.all([first, second]);

            expect(handleFirstConnect).toHaveBeenCalledTimes(1);
        });

        it('should allow initializeClient to sync again after invalidateClient', async () => {
            mockWebSocketClient.isConnected.mockReturnValue(false);

            await manager.initializeClient(mockServerUrl);
            expect(handleFirstConnect).toHaveBeenCalledTimes(1);

            await manager.invalidateClient(mockServerUrl);
            await manager.createClient(mockServerUrl, mockToken);
            mockWebSocketClient.isConnected.mockReturnValue(false);

            await manager.initializeClient(mockServerUrl);
            expect(handleFirstConnect).toHaveBeenCalledTimes(2);
        });

        it('should handle reconnect correctly', async () => {
            const client = manager.getClient(mockServerUrl);
            expect(client).toBeDefined();

            mockCallbacks.reconnect();
            expect(handleReconnect).toHaveBeenCalledWith(mockServerUrl);
        });
    });

    describe('state changes', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
        });

        it('should handle app state changes', () => {
            const mockTimerId = 123;
            jest.spyOn(BackgroundTimer, 'setTimeout').mockReturnValue(mockTimerId);

            // Get the app state callback and simulate background state
            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            expect(BackgroundTimer.setTimeout).toHaveBeenCalled();
            expect(BackgroundTimer.setTimeout).toHaveBeenCalledWith(expect.any(Function), 15000);
        });

        it('background-close timer skips closeAll while a native call is active', () => {
            jest.mocked(hasActiveNativeCall).mockReturnValue(true);

            let capturedCallback: (() => void) | undefined;
            jest.spyOn(BackgroundTimer, 'setTimeout').mockImplementation((cb) => {
                capturedCallback = cb;
                return 1;
            });

            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            jest.clearAllMocks();
            jest.mocked(hasActiveNativeCall).mockReturnValue(true);
            capturedCallback!();

            expect(mockWebSocketClient.close).not.toHaveBeenCalled();
        });

        it('background-close timer closes connections when no native call is active', () => {
            let capturedCallback: (() => void) | undefined;
            jest.spyOn(BackgroundTimer, 'setTimeout').mockImplementation((cb) => {
                capturedCallback = cb;
                return 2;
            });

            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            jest.clearAllMocks();
            jest.mocked(hasActiveNativeCall).mockReturnValue(false);
            capturedCallback!();

            expect(mockWebSocketClient.close).toHaveBeenCalledWith(true);
        });
    });

    describe('scheduleBackgroundCloseIfNeeded', () => {
        beforeEach(async () => {
            await manager.init(mockCredentials);
        });

        it('arms a close timer when backgrounded with no timer running', () => {
            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            jest.clearAllMocks();
            jest.spyOn(BackgroundTimer, 'setTimeout').mockReturnValue(99);

            (manager as any).isBackgroundTimerRunning = false;

            manager.scheduleBackgroundCloseIfNeeded();

            expect(BackgroundTimer.setTimeout).toHaveBeenCalledWith(expect.any(Function), 15000);
        });

        it('is a no-op when the app is in the foreground', () => {
            jest.spyOn(BackgroundTimer, 'setTimeout');

            manager.scheduleBackgroundCloseIfNeeded();

            expect(BackgroundTimer.setTimeout).not.toHaveBeenCalled();
        });

        it('is a no-op when a close timer is already running', () => {
            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            jest.clearAllMocks();
            jest.spyOn(BackgroundTimer, 'setTimeout');

            manager.scheduleBackgroundCloseIfNeeded();

            expect(BackgroundTimer.setTimeout).not.toHaveBeenCalled();
        });

        it('skips closeAll if a new native call started during the grace window', () => {
            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            let capturedCallback: (() => void) | undefined;
            jest.spyOn(BackgroundTimer, 'setTimeout').mockImplementation((cb) => {
                capturedCallback = cb;
                return 99;
            });
            (manager as any).isBackgroundTimerRunning = false;
            manager.scheduleBackgroundCloseIfNeeded();

            jest.clearAllMocks();
            jest.mocked(hasActiveNativeCall).mockReturnValue(true);
            capturedCallback!();

            expect(mockWebSocketClient.close).not.toHaveBeenCalled();
        });

        it('closes connections when no native call is active at timer fire', () => {
            const mockAppStateChange = (AppState.addEventListener as jest.Mock).mock.calls[0][1];
            mockAppStateChange('active');
            mockAppStateChange('background');

            let capturedCallback: (() => void) | undefined;
            jest.spyOn(BackgroundTimer, 'setTimeout').mockImplementation((cb) => {
                capturedCallback = cb;
                return 100;
            });
            (manager as any).isBackgroundTimerRunning = false;
            manager.scheduleBackgroundCloseIfNeeded();

            jest.clearAllMocks();
            jest.mocked(hasActiveNativeCall).mockReturnValue(false);
            capturedCallback!();

            expect(mockWebSocketClient.close).toHaveBeenCalledWith(true);
        });

        it('should handle network state changes', () => {
            const mockNetInfoCallback = (NetInfo.addEventListener as jest.Mock).mock.calls[0][0];
            mockNetInfoCallback({isConnected: false, type: 'none'});

            // Verify that clients are closed when network is disconnected
            expect(manager.getClient(mockServerUrl)?.close).toHaveBeenCalled();
        });
    });

    describe('periodic updates', () => {
        beforeEach(async () => {
            jest.mocked(getCurrentUserId).mockResolvedValue('user1');
            jest.mocked(queryAllUsers).mockImplementation(() => TestHelper.fakeQuery([
                TestHelper.fakeUserModel({id: 'user1'}),
                TestHelper.fakeUserModel({id: 'user2'}),
            ]));
            await manager.init(mockCredentials);
        });

        it('should handle periodic status updates', async () => {
            const client = manager.getClient(mockServerUrl);
            expect(client).toBeDefined();

            // Trigger first connect callback which starts periodic updates
            mockCallbacks.firstConnect();

            // Wait for all promises to resolve
            await new Promise((resolve) => setImmediate(resolve));

            expect(fetchStatusByIds).toHaveBeenCalledWith(mockServerUrl, ['user2']);

            // Close the websocket to stop the periodic updates
            mockCallbacks.close(0);
        });
    });
});