address cr feedback

Author: transphorm

app/tests/utils/proving/provingMachine.generatePayload.test.ts

@@ -70,67 +70,67 @@ const {
   generateTEEInputsDisclose,
 } = require('../../../src/utils/proving/provingInputs');
 
+function setupDefaultStores() {
+  useProvingStore.setState({
+    circuitType: 'register',
+    passportData: { documentCategory: 'passport', mock: false },
+    secret: 'sec',
+    uuid: '123',
+    sharedKey: Buffer.alloc(32, 1),
+    env: 'prod',
+  });
+
+  useSelfAppStore.setState({
+    selfApp: createMockSelfApp(),
+  });
+
+  useProtocolStore.setState({
+    passport: createMockProtocolState(),
+    id_card: createMockProtocolState(),
+  });
+}
+
+function createMockSelfApp() {
+  return {
+    chainID: 42220 as const,
+    userId: 'u',
+    userDefinedData: '0x0',
+    endpointType: 'https' as const,
+    endpoint: 'https://e',
+    scope: 's',
+    sessionId: '',
+    appName: '',
+    logoBase64: '',
+    header: '',
+    userIdType: 'uuid' as const,
+    devMode: false,
+    disclosures: {},
+    version: 1,
+  };
+}
+
+function createMockProtocolState() {
+  return {
+    dsc_tree: 'tree',
+    csca_tree: [['a']],
+    commitment_tree: null,
+    deployed_circuits: null,
+    circuits_dns_mapping: null,
+    alternative_csca: {},
+    fetch_deployed_circuits: jest.fn(() => Promise.resolve()),
+    fetch_circuits_dns_mapping: jest.fn(() => Promise.resolve()),
+    fetch_csca_tree: jest.fn(() => Promise.resolve()),
+    fetch_dsc_tree: jest.fn(() => Promise.resolve()),
+    fetch_identity_tree: jest.fn(() => Promise.resolve()),
+    fetch_alternative_csca: jest.fn(() => Promise.resolve()),
+    fetch_all: jest.fn(() => Promise.resolve()),
+  };
+}
+
 describe('_generatePayload', () => {
   beforeEach(() => {
     jest.clearAllMocks();
-    useProvingStore.setState({
-      circuitType: 'register',
-      passportData: { documentCategory: 'passport', mock: false },
-      secret: 'sec',
-      uuid: '123',
-      sharedKey: Buffer.alloc(32, 1),
-      env: 'prod',
-    });
-    useSelfAppStore.setState({
-      selfApp: {
-        chainID: 42220,
-        userId: 'u',
-        userDefinedData: '0x0',
-        endpointType: 'https',
-        endpoint: 'https://e',
-        scope: 's',
-        sessionId: '',
-        appName: '',
-        logoBase64: '',
-        header: '',
-        userIdType: 'uuid',
-        devMode: false,
-        disclosures: {},
-        version: 1,
-      },
-    });
-    useProtocolStore.setState({
-      passport: {
-        dsc_tree: 'tree',
-        csca_tree: [['a']],
-        commitment_tree: null,
-        deployed_circuits: null,
-        circuits_dns_mapping: null,
-        alternative_csca: {},
-        fetch_deployed_circuits: jest.fn(),
-        fetch_circuits_dns_mapping: jest.fn(),
-        fetch_csca_tree: jest.fn(),
-        fetch_dsc_tree: jest.fn(),
-        fetch_identity_tree: jest.fn(),
-        fetch_alternative_csca: jest.fn(),
-        fetch_all: jest.fn(),
-      },
-      id_card: {
-        commitment_tree: null,
-        dsc_tree: null,
-        csca_tree: null,
-        deployed_circuits: null,
-        circuits_dns_mapping: null,
-        alternative_csca: {},
-        fetch_deployed_circuits: jest.fn(),
-        fetch_circuits_dns_mapping: jest.fn(),
-        fetch_csca_tree: jest.fn(),
-        fetch_dsc_tree: jest.fn(),
-        fetch_identity_tree: jest.fn(),
-        fetch_alternative_csca: jest.fn(),
-        fetch_all: jest.fn(),
-      },
-    } as any);
+    setupDefaultStores();
   });
 
   it('register circuit', async () => {
@@ -152,8 +152,22 @@ describe('_generatePayload', () => {
     useProvingStore.setState({ circuitType: 'dsc' });
     const payload = await useProvingStore.getState()._generatePayload();
     expect(generateTEEInputsDSC).toHaveBeenCalled();
+    expect(generateTEEInputsDSC).toHaveBeenCalledWith(
+      expect.objectContaining({
+        documentCategory: 'passport',
+        mock: false,
+      }),
+      [['a']], // csca_tree
+      'prod', // env
+    );
     expect(useProvingStore.getState().endpointType).toBe('celo');
     expect(payload.params.uuid).toBe('123');
+    expect(payload.params).toEqual({
+      uuid: '123',
+      nonce: [0],
+      cipher_text: [1],
+      auth_tag: [2],
+    });
   });
 
   it('disclose circuit', async () => {
@@ -163,4 +177,48 @@ describe('_generatePayload', () => {
     expect(useProvingStore.getState().endpointType).toBe('https');
     expect(payload.params.uuid).toBe('123');
   });
+
+  it('handles missing passport data', async () => {
+    useProvingStore.setState({ passportData: null });
+    await expect(
+      useProvingStore.getState()._generatePayload(),
+    ).rejects.toThrow();
+  });
+
+  it('handles input generation failure', async () => {
+    // Reset all mocks first
+    jest.clearAllMocks();
+    setupDefaultStores();
+
+    generateTEEInputsRegister.mockImplementation(() => {
+      throw new Error('Input generation failed');
+    });
+    await expect(useProvingStore.getState()._generatePayload()).rejects.toThrow(
+      'Input generation failed',
+    );
+
+    // Restore the mock to its original implementation
+    generateTEEInputsRegister.mockRestore();
+  });
+
+  it('handles encryption failure', async () => {
+    // Reset all mocks first
+    jest.clearAllMocks();
+    setupDefaultStores();
+
+    // Restore all original implementations first
+    generateTEEInputsRegister.mockImplementation(() => ({
+      inputs: { r: 1 },
+      circuitName: 'reg',
+      endpointType: 'celo',
+      endpoint: 'https://reg',
+    }));
+
+    encryptAES256GCM.mockImplementation(() => {
+      throw new Error('Encryption failed');
+    });
+    await expect(useProvingStore.getState()._generatePayload()).rejects.toThrow(
+      'Encryption failed',
+    );
+  });
 });