Add proving machine refactor guardrail tests (#1584)

* Add proving machine refactor tests

* Clarify refactor guardrail test intent

* Add Socket.IO status handler wiring tests for proving store (#1586)

* Add status handler listener tests

* Fix failure status test payload

* Add hello ack uuid test (#1588)

* formatting

* format, agent feedback

* delete mock mock tests

packages/mobile-sdk-alpha/tests/proving/internal/payloadGenerator.test.ts

@@ -0,0 +1,252 @@
+// SPDX-FileCopyrightText: 2025 Social Connect Labs, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+// NOTE: Converts to Apache-2.0 on 2029-06-11 per LICENSE.
+
+import type { SelfClient } from '../../../src';
+import { useProvingStore } from '../../../src/proving/provingMachine';
+import { useProtocolStore } from '../../../src/stores/protocolStore';
+import { useSelfAppStore } from '../../../src/stores/selfAppStore';
+import { actorMock } from '../actorMock';
+
+vitest.mock('xstate', () => {
+  return {
+    createActor: vitest.fn(() => actorMock),
+    createMachine: vitest.fn(),
+    assign: vitest.fn(),
+    send: vitest.fn(),
+    spawn: vitest.fn(),
+    interpret: vitest.fn(),
+    fromPromise: vitest.fn(),
+    fromObservable: vitest.fn(),
+    fromEventObservable: vitest.fn(),
+    fromCallback: vitest.fn(),
+    fromTransition: vitest.fn(),
+    fromReducer: vitest.fn(),
+    fromRef: vitest.fn(),
+  };
+});
+
+vitest.mock('@selfxyz/common/utils/proving', async () => {
+  const actual = await vitest.importActual<typeof import('@selfxyz/common/utils/proving')>(
+    '@selfxyz/common/utils/proving',
+  );
+  return {
+    ...actual,
+    getPayload: vitest.fn(() => ({ payload: true })),
+    encryptAES256GCM: vitest.fn(() => ({
+      nonce: [1],
+      cipher_text: [2],
+      auth_tag: [3],
+    })),
+  };
+});
+
+vitest.mock('@selfxyz/common/utils/circuits/registerInputs', async () => {
+  const actual = (await vitest.importActual('@selfxyz/common/utils/circuits/registerInputs')) as any;
+  return {
+    ...actual,
+    generateTEEInputsRegister: vitest.fn(async () => ({
+      inputs: { reg: true },
+      circuitName: 'register_circuit',
+      endpointType: 'celo',
+      endpoint: 'https://register',
+    })),
+    generateTEEInputsDSC: vitest.fn(() => ({
+      inputs: { dsc: true },
+      circuitName: 'dsc_circuit',
+      endpointType: 'celo',
+      endpoint: 'https://dsc',
+    })),
+    generateTEEInputsDiscloseStateless: vitest.fn(() => ({
+      inputs: { disclose: true },
+      circuitName: 'disclose_circuit',
+      endpointType: 'https',
+      endpoint: 'https://disclose',
+    })),
+  };
+});
+
+describe('payload generator (refactor guardrail via _generatePayload)', () => {
+  const selfClient: SelfClient = {
+    trackEvent: vitest.fn(),
+    emit: vitest.fn(),
+    logProofEvent: vitest.fn(),
+    getPrivateKey: vitest.fn(),
+    getSelfAppState: () => useSelfAppStore.getState(),
+    getProvingState: () => useProvingStore.getState(),
+    getProtocolState: () => useProtocolStore.getState(),
+  } as unknown as SelfClient;
+
+  beforeEach(() => {
+    vitest.clearAllMocks();
+    useSelfAppStore.setState({
+      selfApp: {
+        chainID: 42220,
+        userId: '12345678-1234-1234-1234-123456789abc',
+        userDefinedData: '0x0',
+        selfDefinedData: '',
+        endpointType: 'https',
+        endpoint: 'https://endpoint',
+        scope: 'scope',
+        sessionId: '',
+        appName: '',
+        logoBase64: '',
+        header: '',
+        userIdType: 'uuid',
+        devMode: false,
+        disclosures: {},
+        version: 1,
+        deeplinkCallback: '',
+      },
+    });
+  });
+
+  it('builds a submit request payload with the encrypted payload', async () => {
+    useProvingStore.setState({
+      circuitType: 'register',
+      passportData: { documentCategory: 'passport', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    const payload = await useProvingStore.getState()._generatePayload(selfClient);
+
+    expect(payload).toEqual({
+      jsonrpc: '2.0',
+      method: 'openpassport_submit_request',
+      id: 2,
+      params: {
+        uuid: 'uuid-123',
+        nonce: [1],
+        cipher_text: [2],
+        auth_tag: [3],
+      },
+    });
+  });
+
+  it('throws when dsc is requested for aadhaar documents', async () => {
+    useProvingStore.setState({
+      circuitType: 'dsc',
+      passportData: { documentCategory: 'aadhaar', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await expect(useProvingStore.getState()._generatePayload(selfClient)).rejects.toThrow(
+      'DSC circuit type is not supported for Aadhaar documents',
+    );
+  });
+
+  it('throws when disclose circuit is requested without a SelfApp', async () => {
+    useSelfAppStore.setState({ selfApp: null });
+    useProvingStore.setState({
+      circuitType: 'disclose',
+      passportData: { documentCategory: 'passport', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await expect(useProvingStore.getState()._generatePayload(selfClient)).rejects.toThrow(
+      'SelfApp context not initialized',
+    );
+  });
+
+  it('throws on invalid circuit types', async () => {
+    useProvingStore.setState({
+      circuitType: 'invalid' as any,
+      passportData: { documentCategory: 'passport', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await expect(useProvingStore.getState()._generatePayload(selfClient)).rejects.toThrow(
+      'Invalid circuit type:invalid',
+    );
+  });
+
+  it('uses register_id for id cards', async () => {
+    const { getPayload } = await import('@selfxyz/common/utils/proving');
+
+    useProvingStore.setState({
+      circuitType: 'register',
+      passportData: { documentCategory: 'id_card', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await useProvingStore.getState()._generatePayload(selfClient);
+
+    expect(getPayload).toHaveBeenCalledWith(
+      { reg: true },
+      'register_id',
+      'register_circuit',
+      'celo',
+      'https://register',
+      1,
+      expect.any(String),
+      '',
+    );
+  });
+
+  it('keeps dsc circuit type for passport documents', async () => {
+    const { getPayload } = await import('@selfxyz/common/utils/proving');
+
+    useProvingStore.setState({
+      circuitType: 'dsc',
+      passportData: { documentCategory: 'passport', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await useProvingStore.getState()._generatePayload(selfClient);
+
+    expect(getPayload).toHaveBeenCalledWith(
+      { dsc: true },
+      'dsc',
+      'dsc_circuit',
+      'celo',
+      'https://dsc',
+      1,
+      expect.any(String),
+      '',
+    );
+  });
+
+  it('always uses disclose for disclosure flows', async () => {
+    const { getPayload } = await import('@selfxyz/common/utils/proving');
+
+    useProvingStore.setState({
+      circuitType: 'disclose',
+      passportData: { documentCategory: 'passport', mock: false } as any,
+      secret: 'secret',
+      uuid: 'uuid-123',
+      sharedKey: Buffer.alloc(32, 1),
+      env: 'prod',
+    });
+
+    await useProvingStore.getState()._generatePayload(selfClient);
+
+    expect(getPayload).toHaveBeenCalledWith(
+      { disclose: true },
+      'disclose',
+      'disclose_circuit',
+      'https',
+      'https://disclose',
+      1,
+      expect.any(String),
+      '',
+    );
+  });
+});

packages/mobile-sdk-alpha/tests/proving/internal/statusListener.test.ts

@@ -0,0 +1,172 @@
+// SPDX-FileCopyrightText: 2025 Social Connect Labs, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+// NOTE: Converts to Apache-2.0 on 2029-06-11 per LICENSE.
+
+import { EventEmitter } from 'events';
+import type { Socket } from 'socket.io-client';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { useProvingStore } from '../../../src/proving/provingMachine';
+import { actorMock } from '../actorMock';
+
+vi.mock('socket.io-client');
+vi.mock('../../../src/constants/analytics', () => ({
+  ProofEvents: {
+    SOCKETIO_CONN_STARTED: 'SOCKETIO_CONN_STARTED',
+    SOCKETIO_SUBSCRIBED: 'SOCKETIO_SUBSCRIBED',
+    SOCKETIO_STATUS_RECEIVED: 'SOCKETIO_STATUS_RECEIVED',
+    SOCKETIO_PROOF_FAILURE: 'SOCKETIO_PROOF_FAILURE',
+    SOCKETIO_PROOF_SUCCESS: 'SOCKETIO_PROOF_SUCCESS',
+    REGISTER_COMPLETED: 'REGISTER_COMPLETED',
+  },
+  PassportEvents: {},
+}));
+
+vi.mock('../../../src/proving/internal/logging', () => ({
+  logProofEvent: vi.fn(),
+  createProofContext: vi.fn(() => ({})),
+}));
+vi.mock('@selfxyz/common/utils/proving', () => ({
+  getWSDbRelayerUrl: vi.fn(() => 'ws://test-url'),
+  getPayload: vi.fn(),
+  encryptAES256GCM: vi.fn(),
+  clientKey: {},
+  clientPublicKeyHex: 'test-key',
+  ec: {},
+}));
+
+vi.mock('../../../src/documents/utils', () => ({
+  loadSelectedDocument: vi.fn(() =>
+    Promise.resolve({
+      data: { mockData: true },
+      version: '1.0.0',
+    }),
+  ),
+  hasAnyValidRegisteredDocument: vi.fn(() => Promise.resolve(true)),
+  clearPassportData: vi.fn(),
+  markCurrentDocumentAsRegistered: vi.fn(),
+  reStorePassportDataWithRightCSCA: vi.fn(),
+}));
+
+vi.mock('../../../src/types/events', () => ({
+  SdkEvents: {
+    PASSPORT_DATA_NOT_FOUND: 'PASSPORT_DATA_NOT_FOUND',
+  },
+}));
+
+vi.mock('@selfxyz/common/utils', () => ({
+  getCircuitNameFromPassportData: vi.fn(() => 'register'),
+  getSolidityPackedUserContextData: vi.fn(() => '0x123'),
+}));
+
+vi.mock('@selfxyz/common/utils/attest', () => ({
+  getPublicKey: vi.fn(),
+  verifyAttestation: vi.fn(),
+}));
+
+vi.mock('@selfxyz/common/utils/circuits/registerInputs', () => ({
+  generateTEEInputsDSC: vi.fn(),
+  generateTEEInputsRegister: vi.fn(),
+}));
+
+vi.mock('@selfxyz/common/utils/passports/validate', () => ({
+  checkDocumentSupported: vi.fn(() => Promise.resolve(true)),
+  checkIfPassportDscIsInTree: vi.fn(() => Promise.resolve(true)),
+  isDocumentNullified: vi.fn(() => Promise.resolve(false)),
+  isUserRegistered: vi.fn(() => Promise.resolve(false)),
+  isUserRegisteredWithAlternativeCSCA: vi.fn(() => Promise.resolve(false)),
+}));
+
+vi.mock('xstate', () => ({
+  createActor: vi.fn(() => actorMock),
+  createMachine: vi.fn(() => ({})),
+}));
+
+describe('Socket.IO status handler wiring', () => {
+  const mockSelfClient = {
+    trackEvent: vi.fn(),
+    emit: vi.fn(),
+    getPrivateKey: vi.fn(() => Promise.resolve('mock-private-key')),
+    logProofEvent: vi.fn(),
+    getSelfAppState: () => ({
+      selfApp: {},
+    }),
+    getProtocolState: () => ({
+      isUserLoggedIn: true,
+    }),
+    getProvingState: () => useProvingStore.getState(),
+  } as any;
+
+  let mockSocket: EventEmitter & Partial<Socket>;
+  let socketIoMock: any;
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+
+    useProvingStore.setState({
+      socketConnection: null,
+      error_code: null,
+      reason: null,
+      circuitType: 'register',
+    } as any);
+
+    mockSocket = new EventEmitter() as EventEmitter & Partial<Socket>;
+    vi.spyOn(mockSocket as any, 'emit');
+    mockSocket.disconnect = vi.fn();
+
+    const socketIo = await import('socket.io-client');
+    socketIoMock = vi.mocked(socketIo.default || socketIo);
+    socketIoMock.mockReturnValue(mockSocket);
+
+    const store = useProvingStore.getState();
+    await store.init(mockSelfClient, 'register', true);
+    actorMock.send.mockClear();
+  });
+
+  it('applies success updates and emits PROVE_SUCCESS', async () => {
+    const store = useProvingStore.getState();
+    store._startSocketIOStatusListener('test-uuid', 'https', mockSelfClient);
+
+    await new Promise(resolve => setImmediate(resolve));
+
+    (mockSocket as any).emit('status', { status: 4 });
+
+    const finalState = useProvingStore.getState();
+    expect(finalState.socketConnection).toBe(null);
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_SUCCESS' });
+  });
+
+  it('applies failure updates and emits PROVE_FAILURE', async () => {
+    const store = useProvingStore.getState();
+    store._startSocketIOStatusListener('test-uuid', 'https', mockSelfClient);
+
+    await new Promise(resolve => setImmediate(resolve));
+
+    (mockSocket as any).emit('status', {
+      status: 5,
+      error_code: 'E001',
+      reason: 'TEE failed',
+    });
+
+    const finalState = useProvingStore.getState();
+    expect(finalState.error_code).toBe('E001');
+    expect(finalState.reason).toBe('TEE failed');
+    expect(finalState.socketConnection).toBe(null);
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_FAILURE' });
+  });
+
+  it('emits PROVE_ERROR without updating state for retryable errors', async () => {
+    const store = useProvingStore.getState();
+    store._startSocketIOStatusListener('test-uuid', 'https', mockSelfClient);
+
+    await new Promise(resolve => setImmediate(resolve));
+
+    (mockSocket as any).emit('status', '{"invalid": json}');
+
+    const finalState = useProvingStore.getState();
+    expect(finalState.socketConnection).toBe(mockSocket);
+    expect(finalState.error_code).toBe(null);
+    expect(finalState.reason).toBe(null);
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_ERROR' });
+  });
+});

packages/mobile-sdk-alpha/tests/proving/internal/websocketHandlers.test.ts

@@ -0,0 +1,227 @@
+// SPDX-FileCopyrightText: 2025 Social Connect Labs, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+// NOTE: Converts to Apache-2.0 on 2029-06-11 per LICENSE.
+
+import type { SelfClient } from '../../../src';
+import * as documentUtils from '../../../src/documents/utils';
+import { useProvingStore } from '../../../src/proving/provingMachine';
+import { useProtocolStore } from '../../../src/stores/protocolStore';
+import { useSelfAppStore } from '../../../src/stores/selfAppStore';
+import { actorMock } from '../actorMock';
+
+vitest.mock('uuid', () => ({
+  v4: vitest.fn(() => 'uuid-123'),
+}));
+
+vitest.mock('xstate', () => {
+  return {
+    createActor: vitest.fn(() => actorMock),
+    createMachine: vitest.fn(),
+    assign: vitest.fn(),
+    send: vitest.fn(),
+    spawn: vitest.fn(),
+    interpret: vitest.fn(),
+    fromPromise: vitest.fn(),
+    fromObservable: vitest.fn(),
+    fromEventObservable: vitest.fn(),
+    fromCallback: vitest.fn(),
+    fromTransition: vitest.fn(),
+    fromReducer: vitest.fn(),
+    fromRef: vitest.fn(),
+  };
+});
+
+vitest.mock('@selfxyz/common/utils/attest', () => {
+  return {
+    validatePKIToken: vitest.fn(() => ({
+      userPubkey: Buffer.from('abcd', 'hex'),
+      serverPubkey: 'server-key',
+      imageHash: 'hash',
+      verified: true,
+    })),
+    checkPCR0Mapping: vitest.fn(async () => true),
+  };
+});
+
+vitest.mock('@selfxyz/common/utils/proving', async () => {
+  const actual = await vitest.importActual<typeof import('@selfxyz/common/utils/proving')>(
+    '@selfxyz/common/utils/proving',
+  );
+  return {
+    ...actual,
+    clientPublicKeyHex: 'abcd',
+    clientKey: {
+      derive: vitest.fn(() => ({
+        toArray: () => Array(32).fill(7),
+      })),
+    },
+    ec: {
+      keyFromPublic: vitest.fn(() => ({
+        getPublic: vitest.fn(() => 'server-public'),
+      })),
+    },
+  };
+});
+
+describe('websocket handlers (refactor guardrail via proving store)', () => {
+  const selfClient: SelfClient = {
+    trackEvent: vitest.fn(),
+    emit: vitest.fn(),
+    logProofEvent: vitest.fn(),
+    getPrivateKey: vitest.fn().mockResolvedValue('secret'),
+    getSelfAppState: () => useSelfAppStore.getState(),
+    getProvingState: () => useProvingStore.getState(),
+    getProtocolState: () => useProtocolStore.getState(),
+  } as unknown as SelfClient;
+  let loadSelectedDocumentSpy: any;
+
+  beforeEach(() => {
+    vitest.clearAllMocks();
+    (globalThis as { __DEV__?: boolean }).__DEV__ = true;
+    useSelfAppStore.setState({ selfApp: null, sessionId: null, socket: null });
+    if (!loadSelectedDocumentSpy) {
+      loadSelectedDocumentSpy = vitest.spyOn(documentUtils, 'loadSelectedDocument');
+    }
+    loadSelectedDocumentSpy.mockResolvedValue({
+      data: {
+        documentCategory: 'passport',
+        mock: false,
+        dsc_parsed: { authorityKeyIdentifier: 'aki' },
+      } as any,
+    } as any);
+  });
+
+  it('does nothing when actor is missing or wsConnection is null', () => {
+    useProvingStore.setState({ wsConnection: null } as any);
+
+    useProvingStore.getState()._handleWsOpen(selfClient);
+
+    expect(actorMock.send).not.toHaveBeenCalled();
+  });
+
+  it('does nothing when wsConnection is null even if actor exists', async () => {
+    await useProvingStore.getState().init(selfClient, 'register');
+    actorMock.send.mockClear();
+    useProvingStore.setState({ wsConnection: null } as any);
+
+    useProvingStore.getState()._handleWsOpen(selfClient);
+
+    expect(actorMock.send).not.toHaveBeenCalled();
+  });
+
+  it('sends hello message and stores uuid on open', async () => {
+    const wsConnection = {
+      send: vitest.fn(),
+    } as unknown as WebSocket;
+
+    await useProvingStore.getState().init(selfClient, 'register');
+    useProvingStore.setState({ wsConnection });
+
+    useProvingStore.getState()._handleWsOpen(selfClient);
+
+    expect(useProvingStore.getState().uuid).toBe('uuid-123');
+    const [sentMessage] = (wsConnection.send as unknown as { mock: { calls: string[][] } }).mock.calls[0];
+    const parsedMessage = JSON.parse(sentMessage);
+    expect(parsedMessage).toMatchObject({
+      jsonrpc: '2.0',
+      method: 'openpassport_hello',
+      id: 1,
+      params: {
+        uuid: 'uuid-123',
+        user_pubkey: expect.any(Array),
+      },
+    });
+  });
+
+  it('handles attestation messages by deriving shared key and emitting CONNECT_SUCCESS', async () => {
+    const { clientPublicKeyHex } = await import('@selfxyz/common/utils/proving');
+    const { validatePKIToken } = await import('@selfxyz/common/utils/attest');
+
+    await useProvingStore.getState().init(selfClient, 'register');
+    useProvingStore.setState({ currentState: 'init_tee_connexion' } as any);
+
+    const event = { data: JSON.stringify({ result: { attestation: [1, 2, 3] } }) } as MessageEvent;
+
+    await useProvingStore.getState()._handleWebSocketMessage(event, selfClient);
+
+    expect(clientPublicKeyHex).toBe('abcd');
+    expect(validatePKIToken).toHaveBeenCalled();
+    expect(useProvingStore.getState().sharedKey).toEqual(Buffer.from(Array(32).fill(7)));
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'CONNECT_SUCCESS' });
+  });
+
+  it('starts socket listener on hello ack', async () => {
+    await useProvingStore.getState().init(selfClient, 'register');
+    const startListener = vitest.fn();
+    useProvingStore.setState({
+      endpointType: 'https',
+      uuid: 'uuid-123',
+      _startSocketIOStatusListener: startListener,
+    } as any);
+
+    const event = new MessageEvent('message', {
+      data: JSON.stringify({ id: 2, result: 'status-uuid' }),
+    });
+
+    await useProvingStore.getState()._handleWebSocketMessage(event, selfClient);
+
+    expect(startListener).toHaveBeenCalledWith('status-uuid', 'https', selfClient);
+  });
+
+  it('uses hello ack uuid when it differs from stored uuid', async () => {
+    await useProvingStore.getState().init(selfClient, 'register');
+    const startListener = vitest.fn();
+    useProvingStore.setState({
+      endpointType: 'https',
+      uuid: 'uuid-123',
+      _startSocketIOStatusListener: startListener,
+    } as any);
+
+    const event = new MessageEvent('message', {
+      data: JSON.stringify({ id: 2, result: 'uuid-456' }),
+    });
+
+    await useProvingStore.getState()._handleWebSocketMessage(event, selfClient);
+
+    expect(startListener).toHaveBeenCalledWith('uuid-456', 'https', selfClient);
+  });
+
+  it('emits PROVE_ERROR on websocket error payloads', async () => {
+    await useProvingStore.getState().init(selfClient, 'register');
+
+    const event = new MessageEvent('message', {
+      data: JSON.stringify({ error: 'bad' }),
+    });
+
+    await useProvingStore.getState()._handleWebSocketMessage(event, selfClient);
+
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_ERROR' });
+  });
+
+  it.each([
+    { state: 'init_tee_connexion', expected: 'PROVE_ERROR' },
+    { state: 'proving', expected: 'PROVE_ERROR' },
+    { state: 'listening_for_status', expected: 'PROVE_ERROR' },
+  ])('emits $expected when websocket closes during $state', async ({ state, expected }) => {
+    await useProvingStore.getState().init(selfClient, 'register');
+    useProvingStore.setState({ currentState: state } as any);
+
+    const event = { code: 1000, reason: 'closed' } as CloseEvent;
+
+    useProvingStore.getState()._handleWsClose(event, selfClient);
+
+    expect(actorMock.send).toHaveBeenCalledWith({ type: expected });
+  });
+
+  it.each([
+    { state: 'init_tee_connexion', expected: 'PROVE_ERROR' },
+    { state: 'proving', expected: 'PROVE_ERROR' },
+  ])('emits $expected when websocket errors during $state', async ({ state, expected }) => {
+    await useProvingStore.getState().init(selfClient, 'register');
+    useProvingStore.setState({ currentState: state } as any);
+
+    useProvingStore.getState()._handleWsError(new Event('error'), selfClient);
+
+    expect(actorMock.send).toHaveBeenCalledWith({ type: expected });
+  });
+});

packages/mobile-sdk-alpha/tests/proving/internal/websocketUrlResolver.test.ts

@@ -0,0 +1,205 @@
+// SPDX-FileCopyrightText: 2025 Social Connect Labs, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+// NOTE: Converts to Apache-2.0 on 2029-06-11 per LICENSE.
+
+import type { SelfClient } from '../../../src';
+import * as documentUtils from '../../../src/documents/utils';
+import { useProvingStore } from '../../../src/proving/provingMachine';
+import { useProtocolStore } from '../../../src/stores/protocolStore';
+import { actorMock, emitState } from '../actorMock';
+
+vitest.mock('xstate', () => {
+  return {
+    createActor: vitest.fn(() => actorMock),
+    createMachine: vitest.fn(),
+    assign: vitest.fn(),
+    send: vitest.fn(),
+    spawn: vitest.fn(),
+    interpret: vitest.fn(),
+    fromPromise: vitest.fn(),
+    fromObservable: vitest.fn(),
+    fromEventObservable: vitest.fn(),
+    fromCallback: vitest.fn(),
+    fromTransition: vitest.fn(),
+    fromReducer: vitest.fn(),
+    fromRef: vitest.fn(),
+  };
+});
+
+vitest.mock('@selfxyz/common/utils', async () => {
+  const actual = await vitest.importActual<typeof import('@selfxyz/common/utils')>('@selfxyz/common/utils');
+  return {
+    ...actual,
+    getCircuitNameFromPassportData: vitest.fn(() => 'mock-circuit'),
+  };
+});
+
+describe('websocket URL resolution (refactor guardrail via initTeeConnection)', () => {
+  const wsSend = vitest.fn();
+  const wsAddEventListener = vitest.fn();
+  const wsMock = vitest.fn(() => ({
+    addEventListener: wsAddEventListener,
+    send: wsSend,
+  }));
+  let loadSelectedDocumentSpy: any;
+
+  const makeSelfClient = (): SelfClient =>
+    ({
+      getPrivateKey: vitest.fn().mockResolvedValue('secret'),
+      trackEvent: vitest.fn(),
+      logProofEvent: vitest.fn(),
+      getProvingState: () => useProvingStore.getState(),
+      getProtocolState: () => useProtocolStore.getState(),
+      getSelfAppState: () => ({ selfApp: null }),
+    }) as unknown as SelfClient;
+
+  const setCircuitsMapping = (documentCategory: 'passport' | 'id_card' | 'aadhaar', mapping: any) => {
+    useProtocolStore.setState(state => ({
+      [documentCategory]: {
+        ...state[documentCategory],
+        circuits_dns_mapping: mapping,
+      },
+    }));
+  };
+
+  beforeEach(() => {
+    vitest.restoreAllMocks();
+    vitest.clearAllMocks();
+    global.WebSocket = wsMock as unknown as typeof WebSocket;
+    loadSelectedDocumentSpy = vitest.spyOn(documentUtils, 'loadSelectedDocument');
+  });
+
+  it.each([
+    {
+      label: 'disclose passport -> DISCLOSE',
+      circuitType: 'disclose' as const,
+      documentCategory: 'passport' as const,
+      circuitName: 'disclose',
+      mappingKey: 'DISCLOSE',
+    },
+    {
+      label: 'disclose id_card -> DISCLOSE_ID',
+      circuitType: 'disclose' as const,
+      documentCategory: 'id_card' as const,
+      circuitName: 'disclose',
+      mappingKey: 'DISCLOSE_ID',
+    },
+    {
+      label: 'disclose aadhaar -> DISCLOSE_AADHAAR',
+      circuitType: 'disclose' as const,
+      documentCategory: 'aadhaar' as const,
+      circuitName: 'disclose_aadhaar',
+      mappingKey: 'DISCLOSE_AADHAAR',
+    },
+    {
+      label: 'register passport -> REGISTER',
+      circuitType: 'register' as const,
+      documentCategory: 'passport' as const,
+      circuitName: 'mock-circuit',
+      mappingKey: 'REGISTER',
+    },
+    {
+      label: 'register id_card -> REGISTER_ID',
+      circuitType: 'register' as const,
+      documentCategory: 'id_card' as const,
+      circuitName: 'mock-circuit',
+      mappingKey: 'REGISTER_ID',
+    },
+    {
+      label: 'register aadhaar -> REGISTER_AADHAAR',
+      circuitType: 'register' as const,
+      documentCategory: 'aadhaar' as const,
+      circuitName: 'mock-circuit',
+      mappingKey: 'REGISTER_AADHAAR',
+    },
+    {
+      label: 'dsc passport -> DSC',
+      circuitType: 'dsc' as const,
+      documentCategory: 'passport' as const,
+      circuitName: 'mock-circuit',
+      mappingKey: 'DSC',
+    },
+    {
+      label: 'dsc id_card -> DSC_ID',
+      circuitType: 'dsc' as const,
+      documentCategory: 'id_card' as const,
+      circuitName: 'mock-circuit',
+      mappingKey: 'DSC_ID',
+    },
+  ])('$label resolves expected WebSocket URL', async ({ circuitType, documentCategory, circuitName, mappingKey }) => {
+    const selfClient = makeSelfClient();
+    const wsUrl = `wss://example/${mappingKey}`;
+
+    loadSelectedDocumentSpy.mockResolvedValue({
+      data: {
+        documentCategory,
+        mock: false,
+        dsc_parsed: { authorityKeyIdentifier: 'aki' },
+      } as any,
+    } as any);
+
+    setCircuitsMapping(documentCategory, {
+      [mappingKey]: {
+        [circuitName]: wsUrl,
+      },
+    });
+
+    await useProvingStore.getState().init(selfClient, circuitType);
+
+    const initPromise = useProvingStore.getState().initTeeConnection(selfClient);
+    emitState('ready_to_prove');
+    await initPromise;
+
+    expect(wsMock).toHaveBeenCalledWith(wsUrl);
+  });
+
+  it('throws when mapping is missing for the circuit', async () => {
+    const selfClient = makeSelfClient();
+
+    loadSelectedDocumentSpy.mockResolvedValue({
+      data: {
+        documentCategory: 'passport',
+        mock: false,
+        dsc_parsed: { authorityKeyIdentifier: 'aki' },
+      } as any,
+    } as any);
+
+    setCircuitsMapping('passport', {
+      REGISTER: {
+        other: 'wss://missing',
+      },
+    });
+
+    await useProvingStore.getState().init(selfClient, 'register');
+
+    await expect(useProvingStore.getState().initTeeConnection(selfClient)).rejects.toThrow(
+      'No WebSocket URL available for TEE connection',
+    );
+  });
+
+  it('throws for unsupported document categories', async () => {
+    const selfClient = makeSelfClient();
+    const invalidCategory = 'driver_license';
+
+    loadSelectedDocumentSpy.mockResolvedValue({
+      data: {
+        documentCategory: invalidCategory,
+        mock: false,
+        dsc_parsed: { authorityKeyIdentifier: 'aki' },
+      } as any,
+    } as any);
+
+    useProtocolStore.setState(state => ({
+      ...(state as any),
+      [invalidCategory]: {
+        circuits_dns_mapping: {},
+      },
+    }));
+
+    await useProvingStore.getState().init(selfClient, 'disclose');
+
+    await expect(useProvingStore.getState().initTeeConnection(selfClient)).rejects.toThrow(
+      'Unsupported document category for disclose: driver_license',
+    );
+  });
+});