feat: add uninstall validation, tests

alchemyplatform · Dec 14, 2024 · 9e5c594 · 9e5c594
1 parent 5bd7bcd
commit 9e5c594
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 37 deletions.
diff --git a/account-kit/smart-contracts/src/ma-v2/actions/common/types.ts b/account-kit/smart-contracts/src/ma-v2/actions/common/types.ts
@@ -1,12 +1,12 @@
 import type { Address, Hex } from "viem";
 
 export type ModuleEntity = {
-  address: Address;
+  moduleAddress: Address;
   entityId: number;
 };
 
 export type ValidationConfig = {
-  address: Address;
+  moduleAddress: Address;
   entityId: number; // uint32
   isGlobal: boolean;
   isSignatureValidation: boolean;

diff --git a/account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts b/account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts
@@ -26,5 +26,5 @@ export function serializeHookConfig(config: HookConfig): Hex {
 }
 
 export function serializeModuleEntity(config: ModuleEntity): Hex {
-  return concatHex([config.address, toHex(config.entityId, { size: 4 })]);
+  return concatHex([config.moduleAddress, toHex(config.entityId, { size: 4 })]);
 }
diff --git a/account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts b/account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts
@@ -24,6 +24,7 @@ import type { HookConfig, ValidationConfig } from "../common/types.js";
 import {
   serializeValidationConfig,
   serializeHookConfig,
+  serializeModuleEntity,
 } from "../common/utils.js";
 
 export type InstallValidationParams<
@@ -118,39 +119,40 @@ export const installValidationActions: <
   },
 
   uninstallValidation: async ({
-    // moduleAddress, entityId, uninstallData, hookUninstallDatas,
+    moduleAddress,
+    entityId,
+    uninstallData,
+    hookUninstallDatas,
     account = client.account,
     overrides,
   }) => {
-    // if (!account) {
-    //   throw new AccountNotFoundError();
-    // }
-
-    // if (!isSmartAccountClient(client)) {
-    //   throw new IncompatibleClientError(
-    //     "SmartAccountClient",
-    //     "uninstallValidation",
-    //     client
-    //   );
-    // }
+    if (!account) {
+      throw new AccountNotFoundError();
+    }
 
-    // const { moduleAddress, entityId, uninstallData, hookUninstallDatas } = args;
+    if (!isSmartAccountClient(client)) {
+      throw new IncompatibleClientError(
+        "SmartAccountClient",
+        "uninstallValidation",
+        client
+      );
+    }
 
-    // const callData = encodeFunctionData({
-    //   abi: semiModularAccountBytecodeAbi,
-    //   functionName: "uninstallValidation",
-    //   args: [
-    //     serializeModuleEntity({
-    //       moduleAddress,
-    //       entityId,
-    //     }),
-    //     uninstallData,
-    //     hookUninstallDatas,
-    //   ],
-    // });
+    const callData = encodeFunctionData({
+      abi: semiModularAccountBytecodeAbi,
+      functionName: "uninstallValidation",
+      args: [
+        serializeModuleEntity({
+          moduleAddress,
+          entityId,
+        }),
+        uninstallData,
+        hookUninstallDatas,
+      ],
+    });
 
     return client.sendUserOperation({
-      uo: "0x",
+      uo: callData,
       account,
       overrides,
     });

diff --git a/account-kit/smart-contracts/src/ma-v2/client/client.test.ts b/account-kit/smart-contracts/src/ma-v2/client/client.test.ts
@@ -72,9 +72,9 @@ describe("MA v2 Tests", async () => {
       value: parseEther("2"),
     });
 
-    const result = await provider.installValidation({
+    let result = await provider.installValidation({
       validationConfig: {
-        address: addresses.singleSignerValidationModule,
+        moduleAddress: addresses.singleSignerValidationModule,
         entityId: 1,
         isGlobal: true,
         isSignatureValidation: true,
@@ -93,34 +93,104 @@ describe("MA v2 Tests", async () => {
 
     const startingAddressBalance = await getTargetBalance();
 
-    // connect session key
-    provider.signer = sessionKey;
-    provider.entityId = 1n;
+    // connect session key and send tx with session key
+    let sessionKeyClient = await createSMAV2AccountClient({
+      chain: instance.chain,
+      signer: sessionKey,
+      transport: custom(instance.getClient()),
+      accountAddress: provider.getAddress(),
+      entityId: 1n,
+      isGlobalValidation: true,
+    });
 
-    const result2 = await provider.sendUserOperation({
+    result = await sessionKeyClient.sendUserOperation({
       uo: {
         target: target,
         value: sendAmount,
         data: "0x",
       },
     });
 
-    txnHash = provider.waitForUserOperationTransaction(result);
+    txnHash = sessionKeyClient.waitForUserOperationTransaction(result);
     await expect(txnHash).resolves.not.toThrowError();
-
     await expect(await getTargetBalance()).toEqual(
       startingAddressBalance + sendAmount
     );
   });
 
+  it("uninstalls a session key", async () => {
+    let provider = (await givenConnectedProvider({ signer })).extend(
+      installValidationActions
+    );
+
+    await setBalance(client, {
+      address: provider.getAddress(),
+      value: parseEther("2"),
+    });
+
+    let result = await provider.installValidation({
+      validationConfig: {
+        moduleAddress: addresses.singleSignerValidationModule,
+        entityId: 1,
+        isGlobal: true,
+        isSignatureValidation: true,
+        isUserOpValidation: true,
+      },
+      selectors: [],
+      installData: SingleSignerValidationModule.encodeOnInstallData({
+        entityId: 1,
+        signer: await sessionKey.getAddress(),
+      }),
+      hooks: [],
+    });
+
+    let txnHash = provider.waitForUserOperationTransaction(result);
+    await expect(txnHash).resolves.not.toThrowError();
+
+    result = await provider.uninstallValidation({
+      moduleAddress: addresses.singleSignerValidationModule,
+      entityId: 1,
+      uninstallData: SingleSignerValidationModule.encodeOnUninstallData({
+        entityId: 1,
+      }),
+      hookUninstallDatas: [],
+    });
+
+    txnHash = provider.waitForUserOperationTransaction(result);
+    await expect(txnHash).resolves.not.toThrowError();
+
+    // connect session key and send tx with session key
+    let sessionKeyClient = await createSMAV2AccountClient({
+      chain: instance.chain,
+      signer: sessionKey,
+      transport: custom(instance.getClient()),
+      accountAddress: provider.getAddress(),
+      entityId: 1n,
+      isGlobalValidation: true,
+    });
+
+    result = sessionKeyClient.sendUserOperation({
+      uo: {
+        target: target,
+        value: sendAmount,
+        data: "0x",
+      },
+    });
+
+    await expect(result).rejects.toThrowError();
+  });
+
   const givenConnectedProvider = async ({
     signer,
+    accountAddress,
   }: {
     signer: SmartAccountSigner;
+    accountAddress?: `0x${string}`;
   }) =>
     createSMAV2AccountClient({
       chain: instance.chain,
       signer,
+      accountAddress,
       transport: custom(instance.getClient()),
     });
 });