test(wrapped-keys): LIT-3956 - Add testFailStoreEncryptedKeyBatchIsAtomic test case

Author: MaximusHaximus

local-tests/test.ts

@@ -105,6 +105,7 @@ import { testExportWrappedKey } from './tests/wrapped-keys/testExportWrappedKey'
 import { testSignMessageWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignMessageWithSolanaEncryptedKey';
 import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignTransactionWithSolanaEncryptedKey';
 import { testBatchGeneratePrivateKeys } from './tests/wrapped-keys/testBatchGeneratePrivateKeys';
+import { testFailBatchGeneratePrivateKeysAtomic } from './tests/wrapped-keys/testFailStoreEncryptedKeyBatchIsAtomic';
 
 import { setLitActionsCodeToLocal } from './tests/wrapped-keys/util';
 import { testUseEoaSessionSigsToRequestSingleResponse } from './tests/testUseEoaSessionSigsToRequestSingleResponse';
@@ -151,6 +152,7 @@ setLitActionsCodeToLocal();
     testFailEthereumSignTransactionWrappedKeyWithMissingParam,
     testFailEthereumSignTransactionWrappedKeyWithInvalidParam,
     testFailEthereumSignTransactionWrappedKeyInvalidDecryption,
+    testFailBatchGeneratePrivateKeysAtomic,
 
     // -- import wrapped keys
     testFailImportWrappedKeysWithSamePrivateKey,

local-tests/tests/wrapped-keys/testBatchGeneratePrivateKeys.ts

@@ -7,7 +7,7 @@ import bs58 from 'bs58';
 import { ethers } from 'ethers';
 import { BatchGeneratePrivateKeysActionResult } from '../../../packages/wrapped-keys/src/lib/types';
 
-const { batchGeneratePrivateKeys, exportPrivateKey } = api;
+const { batchGeneratePrivateKeys } = api;
 
 async function verifySolanaSignature(
   solanaResult: BatchGeneratePrivateKeysActionResult,

local-tests/tests/wrapped-keys/testFailStoreEncryptedKeyBatchIsAtomic.ts

@@ -0,0 +1,134 @@
+import { log } from '@lit-protocol/misc';
+import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
+import { api } from '@lit-protocol/wrapped-keys';
+import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
+import nacl from 'tweetnacl';
+import bs58 from 'bs58';
+import { ethers } from 'ethers';
+import { BatchGeneratePrivateKeysActionResult } from '../../../packages/wrapped-keys/src/lib/types';
+import { batchGenerateKeysWithLitAction } from '../../../packages/wrapped-keys/src/lib/lit-actions-client';
+import {
+  getFirstSessionSig,
+  getPkpAccessControlCondition,
+  getPkpAddressFromSessionSig,
+} from '../../../packages/wrapped-keys/src/lib/utils';
+import { getLitActionCodeOrCidCommon } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/utils';
+import { getKeyTypeFromNetwork } from '../../../packages/wrapped-keys/src/lib/api/utils';
+import { listEncryptedKeyMetadata } from '../../../packages/wrapped-keys/src/lib/api';
+
+const { batchGeneratePrivateKeys, storeEncryptedKeyBatch } = api;
+
+/**
+ * Test Commands:
+ * ✅ NETWORK=datil-dev yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ * ✅ NETWORK=datil-test yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ * ✅ NETWORK=localchain yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ */
+export const testFailBatchGeneratePrivateKeysAtomic = async (
+  devEnv: TinnyEnvironment
+) => {
+  const alice = await devEnv.createRandomPerson();
+
+  try {
+    const pkpSessionSigsSigning = await getPkpSessionSigs(
+      devEnv,
+      alice,
+      null,
+      new Date(Date.now() + 1000 * 60 * 10).toISOString()
+    ); // 10 mins expiry
+
+    const solanaMessageToSign = 'This is a test solana message';
+    const evmMessageToSign = 'This is a test evm message';
+
+    const sessionSig = getFirstSessionSig(pkpSessionSigsSigning);
+    const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+
+    const allowPkpAddressToDecrypt = getPkpAccessControlCondition(pkpAddress);
+
+    const { litActionCode, litActionIpfsCid } = getLitActionCodeOrCidCommon(
+      'batchGenerateEncryptedKeys'
+    );
+
+    const actionResults = await batchGenerateKeysWithLitAction({
+      litNodeClient: devEnv.litNodeClient,
+      litActionIpfsCid: litActionCode ? undefined : litActionIpfsCid,
+      litActionCode: litActionCode ? litActionCode : undefined,
+      accessControlConditions: [allowPkpAddressToDecrypt],
+      actions: [
+        {
+          network: 'evm',
+          signMessageParams: { messageToSign: evmMessageToSign },
+          generateKeyParams: { memo: 'Test evm key' },
+        },
+        {
+          network: 'solana',
+          signMessageParams: { messageToSign: solanaMessageToSign },
+          generateKeyParams: { memo: 'Test solana key' },
+        },
+      ],
+      pkpSessionSigs: pkpSessionSigsSigning,
+    });
+
+    const keyParamsBatch = actionResults.map((keyData) => {
+      const { generateEncryptedPrivateKey, network } = keyData;
+      return {
+        ...generateEncryptedPrivateKey,
+        keyType: getKeyTypeFromNetwork(network),
+      };
+    });
+
+    // Intentional failure to persist due to missing publicKey
+    delete keyParamsBatch[0].publicKey;
+
+    try {
+      await storeEncryptedKeyBatch({
+        pkpSessionSigs: pkpSessionSigsSigning,
+        litNodeClient: devEnv.litNodeClient,
+        keyBatch: keyParamsBatch,
+      });
+
+      throw new Error(
+        'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
+      );
+    } catch (err) {
+      // We expect `storeEncryptedKeyBatch` to fail w/ a specific error
+      if (
+        err.message.includes(
+          'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
+        ) ||
+        !err.message.includes(
+          'keyParamsBatch[0]: Missing "publicKey" parameter in request'
+        )
+      ) {
+        throw err;
+      }
+
+      try {
+        const keys = await listEncryptedKeyMetadata({
+          litNodeClient: devEnv.litNodeClient,
+          pkpSessionSigs: pkpSessionSigsSigning,
+        });
+
+        console.error(
+          'Got a value back we shouldnt have from listEncryptedKeyMetadata()',
+          keys
+        );
+
+        throw new Error(
+          'Expected `listEncryptedKeyMetadata() to fail, but it didnt!`'
+        );
+      } catch (err) {
+        if (err.message.includes('No keys exist for pkpAddress')) {
+          log('✅ testFailBatchGeneratePrivateKeysAtomic');
+        } else {
+          throw err;
+        }
+      }
+    }
+  } catch (err) {
+    console.log(err.message, err, err.stack);
+    throw err;
+  } finally {
+    devEnv.releasePrivateKeyFromUser(alice);
+  }
+};