@cofhe/hardhat-plugin 0.2.0 → 0.3.0

package/src/index.ts

@@ -6,32 +6,25 @@ import { extendConfig, extendEnvironment, task, types } from 'hardhat/config';
 import { TASK_TEST, TASK_NODE } from 'hardhat/builtin-tasks/task-names';
 import { HardhatEthersSigner } from '@nomicfoundation/hardhat-ethers/signers';
 import {
+  type CofheClient,
+  type CofheConfig,
+  type CofheInputConfig,
   MOCKS_ZK_VERIFIER_SIGNER_ADDRESS,
-  type CofhesdkClient,
-  type CofhesdkConfig,
-  type CofhesdkInputConfig,
 } from '@cofhe/sdk';
-import { createCofhesdkClient, createCofhesdkConfig } from '@cofhe/sdk/node';
+import { createCofheClient, createCofheConfig } from '@cofhe/sdk/node';
 import { HardhatSignerAdapter } from '@cofhe/sdk/adapters';
 
 import { localcofheFundAccount } from './fund.js';
 import { TASK_COFHE_MOCKS_DEPLOY, TASK_COFHE_MOCKS_SET_LOG_OPS, TASK_COFHE_USE_FAUCET } from './consts.js';
 import { deployMocks, type DeployMocksArgs } from './deploy.js';
 import { mock_setLoggingEnabled, mock_withLogs } from './logging.js';
-import { mock_expectPlaintext } from './utils.js';
+import { getFixedMockContract, mock_expectPlaintext } from './utils.js';
 import { mock_getPlaintext } from './utils.js';
 import type { Contract } from 'ethers';
-import {
-  MockACLArtifact,
-  MockQueryDecrypterArtifact,
-  MockTaskManagerArtifact,
-  MockZkVerifierArtifact,
-  TestBedArtifact,
-} from '@cofhe/mock-contracts';
 import { hardhat } from '@cofhe/sdk/chains';
-export {
+import {
   MockACLArtifact,
-  MockQueryDecrypterArtifact,
+  MockThresholdNetworkArtifact,
   MockTaskManagerArtifact,
   MockZkVerifierArtifact,
   TestBedArtifact,
@@ -43,7 +36,7 @@ export {
  */
 declare module 'hardhat/types/config' {
   interface HardhatUserConfig {
-    cofhesdk?: {
+    cofhe?: {
       /** Whether to log mock operations (default: true) */
       logMocks?: boolean;
       /** Whether to show gas usage warnings for mock operations (default: true) */
@@ -52,7 +45,7 @@ declare module 'hardhat/types/config' {
   }
 
   interface HardhatConfig {
-    cofhesdk: {
+    cofhe: {
       /** Whether to log mock operations (default: true) */
       logMocks: boolean;
       /** Whether to show gas usage warnings for mock operations (default: true) */
@@ -112,9 +105,9 @@ extendConfig((config, userConfig) => {
   }
 
   // Add cofhe config
-  config.cofhesdk = {
-    logMocks: userConfig.cofhesdk?.logMocks ?? true,
-    gasWarning: userConfig.cofhesdk?.gasWarning ?? true,
+  config.cofhe = {
+    logMocks: userConfig.cofhe?.logMocks ?? true,
+    gasWarning: userConfig.cofhe?.gasWarning ?? true,
   };
 });
 
@@ -155,24 +148,41 @@ task(TASK_COFHE_MOCKS_DEPLOY, 'Deploys the mock contracts on the Hardhat network
   .setAction(async ({ deployTestBed, silent }: DeployMocksArgs, hre) => {
     await deployMocks(hre, {
       deployTestBed: deployTestBed ?? true,
-      gasWarning: hre.config.cofhesdk.gasWarning ?? true,
+      gasWarning: hre.config.cofhe.gasWarning ?? true,
       silent: silent ?? false,
     });
   });
 
+// Hardhat plugin auto-deploys mocks for every hardhat test run by overriding TASK_TEST and calling deployMocks(...) before runSuper()
 task(TASK_TEST, 'Deploy mock contracts on hardhat').setAction(async ({}, hre, runSuper) => {
-  await deployMocks(hre, {
-    deployTestBed: true,
-    gasWarning: hre.config.cofhesdk.gasWarning ?? true,
-  });
+  const skipAutoDeploy = (() => {
+    const raw = process.env.COFHE_SKIP_MOCKS_DEPLOY ?? '';
+    const normalized = raw.trim().toLowerCase();
+    return normalized === '1' || normalized === 'true' || normalized === 'yes';
+  })();
+
+  if (!skipAutoDeploy) {
+    await deployMocks(hre, {
+      deployTestBed: true,
+      gasWarning: hre.config.cofhe.gasWarning ?? true,
+    });
+  }
   return runSuper();
 });
 
 task(TASK_NODE, 'Deploy mock contracts on hardhat').setAction(async ({}, hre, runSuper) => {
-  await deployMocks(hre, {
-    deployTestBed: true,
-    gasWarning: hre.config.cofhesdk.gasWarning ?? true,
-  });
+  const skipAutoDeploy = (() => {
+    const raw = process.env.COFHE_SKIP_MOCKS_DEPLOY ?? '';
+    const normalized = raw.trim().toLowerCase();
+    return normalized === '1' || normalized === 'true' || normalized === 'yes';
+  })();
+
+  if (!skipAutoDeploy) {
+    await deployMocks(hre, {
+      deployTestBed: true,
+      gasWarning: hre.config.cofhe.gasWarning ?? true,
+    });
+  }
   return runSuper();
 });
 
@@ -186,6 +196,7 @@ task(TASK_COFHE_MOCKS_SET_LOG_OPS, 'Set logging for the Mock CoFHE contracts')
 
 // MOCK UTILS
 
+export * from './consts.js';
 export * from './utils.js';
 export * from './fund.js';
 export * from './logging.js';
@@ -197,21 +208,21 @@ export * from './deploy.js';
  */
 declare module 'hardhat/types/runtime' {
   export interface HardhatRuntimeEnvironment {
-    cofhesdk: {
+    cofhe: {
       /**
-       * Create a CoFHE SDK configuration for use with cofhesdk.createCofhesdkClient(...)
-       * @param {CofhesdkInputConfig} config - The CoFHE SDK input configuration
-       * @returns {CofhesdkConfig} The CoFHE SDK configuration
+       * Create a CoFHE configuration for use with hre.cofhe.createClient(...)
+       * @param {CofheInputConfig} config - The CoFHE input configuration
+       * @returns {CofheConfig} The CoFHE configuration
        */
-      createCofhesdkConfig: (config: CofhesdkInputConfig) => Promise<CofhesdkConfig>;
+      createConfig: (config: CofheInputConfig) => Promise<CofheConfig>;
       /**
-       * Create a CoFHE SDK client instance
-       * @param {CofhesdkConfig} config - The CoFHE SDK configuration (use createCofhesdkConfig to create with Node.js defaults)
-       * @returns {Promise<CofhesdkClient>} The CoFHE SDK client instance
+       * Create a CoFHE client instance
+       * @param {CofheConfig} config - The CoFHE configuration (use createCofheConfig to create with Node.js defaults)
+       * @returns {Promise<CofheClient>} The CoFHE client instance
        */
-      createCofhesdkClient: (config: CofhesdkConfig) => CofhesdkClient;
+      createClient: (config: CofheConfig) => CofheClient;
       /**
-       * Create viem clients from a Hardhat ethers signer, to be used with `cofhesdkClient.connect(...)`
+       * Create viem clients from a Hardhat ethers signer, to be used with `cofheClient.connect(...)`
        * @param {HardhatEthersSigner} signer - The Hardhat ethers signer to use
        * @returns {Promise<{ publicClient: PublicClient; walletClient: WalletClient }>} The viem clients
        */
@@ -219,20 +230,20 @@ declare module 'hardhat/types/runtime' {
         signer: HardhatEthersSigner
       ) => Promise<{ publicClient: PublicClient; walletClient: WalletClient }>;
       /**
-       * Connect a CoFHE SDK client with a Hardhat ethers signer
-       * @param {CofhesdkClient} client - The CoFHE SDK client to connect
+       * Connect a CoFHE client with a Hardhat ethers signer
+       * @param {CofheClient} client - The CoFHE client to connect
        * @param {HardhatEthersSigner} signer - The Hardhat ethers signer to use
        * @returns {Promise<void>}
        */
-      connectWithHardhatSigner: (client: CofhesdkClient, signer: HardhatEthersSigner) => Promise<void>;
+      connectWithHardhatSigner: (client: CofheClient, signer: HardhatEthersSigner) => Promise<void>;
       /**
        * Create and connect to a batteries included client.
        * Also generates a self-usage a permit for the signer.
-       * If customization is needed, use createCofhesdkClient and connectWithHardhatSigner.
+       * If customization is needed, use createCofheClient and connectWithHardhatSigner.
        * @param {HardhatEthersSigner} signer - The Hardhat ethers signer to use (optional - defaults to first signer)
-       * @returns {Promise<CofhesdkClient>} The CoFHE SDK client instance
+       * @returns {Promise<CofheClient>} The CoFHE client instance
        */
-      createBatteriesIncludedCofhesdkClient: (signer?: HardhatEthersSigner) => Promise<CofhesdkClient>;
+      createClientWithBatteries: (signer?: HardhatEthersSigner) => Promise<CofheClient>;
 
       mocks: {
         /**
@@ -245,7 +256,7 @@ declare module 'hardhat/types/runtime' {
          * Example usage:
          *
          * ```ts
-         * await hre.cofhesdk.mocks.withLogs("counter.increment()", async () => {
+         * await hre.cofhe.mocks.withLogs("counter.increment()", async () => {
          *   await counter.increment();
          * });
          * ```
@@ -319,10 +330,10 @@ declare module 'hardhat/types/runtime' {
         getMockACL: () => Promise<Contract>;
 
         /**
-         * Get the MockQueryDecrypter contract
-         * @returns {Promise<Contract>} The MockQueryDecrypter contract
+         * Get the MockThresholdNetwork contract
+         * @returns {Promise<Contract>} The MockThresholdNetwork contract
          */
-        getMockQueryDecrypter: () => Promise<Contract>;
+        getMockThresholdNetwork: () => Promise<Contract>;
 
         /**
          * Get the MockZkVerifier contract
@@ -340,53 +351,69 @@ declare module 'hardhat/types/runtime' {
   }
 }
 
+/**
+ * Builds the mocks config for the hardhat plugin.
+ * Defaults `encryptDelay` to `0` so tests run without artificial wait times,
+ * unless the user has explicitly provided a value.
+ */
+export function buildHardhatPluginMocksConfig(
+  mocksConfig: CofheInputConfig['mocks']
+): NonNullable<CofheInputConfig['mocks']> {
+  return {
+    ...mocksConfig,
+    encryptDelay: mocksConfig?.encryptDelay ?? 0,
+  };
+}
+
 extendEnvironment((hre) => {
-  hre.cofhesdk = {
-    createCofhesdkConfig: async (config: CofhesdkInputConfig) => {
+  hre.cofhe = {
+    createConfig: async (config: CofheInputConfig) => {
       // Create zkv wallet client
       // This wallet interacts with the MockZkVerifier contract so that the user's connected wallet doesn't have to
       const zkvHhSigner = await hre.ethers.getImpersonatedSigner(MOCKS_ZK_VERIFIER_SIGNER_ADDRESS);
       const { walletClient: zkvWalletClient } = await HardhatSignerAdapter(zkvHhSigner);
 
       // Inject zkv wallet client into config
+      // Set encryptDelay to 0 on hardhat to avoid waiting for delays during tests
       const configWithZkvWalletClient = {
         environment: 'hardhat' as const,
         ...config,
+        mocks: buildHardhatPluginMocksConfig(config.mocks),
         _internal: {
           ...config._internal,
           zkvWalletClient,
         },
       };
 
-      return createCofhesdkConfig(configWithZkvWalletClient);
+      return createCofheConfig(configWithZkvWalletClient);
     },
-    createCofhesdkClient: (config: CofhesdkConfig) => {
-      return createCofhesdkClient(config);
+    createClient: (config: CofheConfig) => {
+      return createCofheClient(config);
     },
     hardhatSignerAdapter: async (signer: HardhatEthersSigner) => {
       return HardhatSignerAdapter(signer);
     },
-    connectWithHardhatSigner: async (client: CofhesdkClient, signer: HardhatEthersSigner) => {
+    connectWithHardhatSigner: async (client: CofheClient, signer: HardhatEthersSigner) => {
       const { publicClient, walletClient } = await HardhatSignerAdapter(signer);
       return client.connect(publicClient, walletClient);
     },
-    createBatteriesIncludedCofhesdkClient: async (signer?: HardhatEthersSigner) => {
+    createClientWithBatteries: async (signer?: HardhatEthersSigner) => {
       // Get signer if not provided
       if (!signer) {
         [signer] = await hre.ethers.getSigners();
       }
 
       // Create config
-      const config = await hre.cofhesdk.createCofhesdkConfig({
+      const config = await hre.cofhe.createConfig({
         environment: 'hardhat',
         supportedChains: [hardhat],
       });
 
       // Create client
-      const client = hre.cofhesdk.createCofhesdkClient(config);
+      const client = hre.cofhe.createClient(config);
 
       // Connect client
-      await hre.cofhesdk.connectWithHardhatSigner(client, signer);
+      await hre.cofhe.connectWithHardhatSigner(client, signer);
 
       // Create self-usage permit
       await client.permits.createSelf({
@@ -417,23 +444,15 @@ extendEnvironment((hre) => {
         const [signer] = await hre.ethers.getSigners();
         return mock_expectPlaintext(signer.provider, ctHash, expectedValue);
       },
-      getMockTaskManager: async () => {
-        return await hre.ethers.getContractAt(MockTaskManagerArtifact.abi, MockTaskManagerArtifact.fixedAddress);
-      },
+      getMockTaskManager: async () => getFixedMockContract(hre, MockTaskManagerArtifact),
       getMockACL: async () => {
-        const tm = await hre.ethers.getContractAt(MockTaskManagerArtifact.abi, MockTaskManagerArtifact.fixedAddress);
-        const aclAddress = await tm.acl();
-        return await hre.ethers.getContractAt(MockACLArtifact.abi, aclAddress);
-      },
-      getMockQueryDecrypter: async () => {
-        return await hre.ethers.getContractAt(MockQueryDecrypterArtifact.abi, MockQueryDecrypterArtifact.fixedAddress);
-      },
-      getMockZkVerifier: async () => {
-        return await hre.ethers.getContractAt(MockZkVerifierArtifact.abi, MockZkVerifierArtifact.fixedAddress);
-      },
-      getTestBed: async () => {
-        return await hre.ethers.getContractAt(TestBedArtifact.abi, TestBedArtifact.fixedAddress);
+        const taskManager = await getFixedMockContract(hre, MockTaskManagerArtifact);
+        const aclAddress = await taskManager.acl();
+        return hre.ethers.getContractAt(MockACLArtifact.abi, aclAddress);
       },
+      getMockThresholdNetwork: async () => getFixedMockContract(hre, MockThresholdNetworkArtifact),
+      getMockZkVerifier: async () => getFixedMockContract(hre, MockZkVerifierArtifact),
+      getTestBed: async () => getFixedMockContract(hre, TestBedArtifact),
     },
   };
 });

package/src/logging.ts

@@ -1,23 +1,17 @@
 import chalk from 'chalk';
 import { type HardhatRuntimeEnvironment } from 'hardhat/types';
-import { TASK_MANAGER_ADDRESS } from './consts';
+import { getFixedMockContract } from './utils';
 import { MockTaskManagerArtifact } from '@cofhe/mock-contracts';
 
-const getDeployedMockTaskManager = async (hre: HardhatRuntimeEnvironment) => {
-  // Fetch the deployed MockTaskManager
-  const taskManager = await hre.ethers.getContractAt(MockTaskManagerArtifact.abi, TASK_MANAGER_ADDRESS);
-
-  return taskManager;
-};
-
-const getLoggingEnabled = async (hre: HardhatRuntimeEnvironment) => {
-  const taskManager = await getDeployedMockTaskManager(hre);
-  return await taskManager.logOps();
+const getLoggingEnabled = async (hre: HardhatRuntimeEnvironment): Promise<boolean> => {
+  const taskManager = await getFixedMockContract(hre, MockTaskManagerArtifact);
+  return taskManager.logOps();
 };
 
 const setLoggingEnabled = async (hre: HardhatRuntimeEnvironment, enabled: boolean) => {
-  const taskManager = await getDeployedMockTaskManager(hre);
-  await taskManager.setLogOps(enabled);
+  const taskManager = await getFixedMockContract(hre, MockTaskManagerArtifact);
+  const tx = await taskManager.setLogOps(enabled);
+  await tx.wait();
 };
 
 // prettier-ignore

package/src/utils.ts

@@ -1,7 +1,38 @@
-import { TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from './consts.js';
+import { TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from '@cofhe/sdk';
 import { expect } from 'chai';
-import { ethers } from 'ethers';
+import { Contract, ethers } from 'ethers';
 import { type HardhatEthersProvider } from '@nomicfoundation/hardhat-ethers/internal/hardhat-ethers-provider';
+import type { MockArtifact } from '@cofhe/mock-contracts';
+import type { HardhatRuntimeEnvironment } from 'hardhat/types';
+
+// Deployment utils
+
+/// Deploys a mock contract from a pre-built artifact from the mock-contracts package
+/// If the mock contract should be deployed to a fixed address, `hardhat_setCode` op is used to set the code at the fixed address
+/// Otherwise, we deploy the contract using ethers.js to a non-fixed address
+export const deployMockContractFromArtifact = async (hre: HardhatRuntimeEnvironment, artifact: MockArtifact) => {
+  // Use hardhat_setCode to deploy to fixed address
+  if (artifact.isFixed) {
+    await hre.network.provider.send('hardhat_setCode', [artifact.fixedAddress, artifact.deployedBytecode]);
+    return getFixedMockContract(hre, artifact);
+  }
+
+  // Use ethers.js to deploy to variable address
+  const [signer] = await hre.ethers.getSigners();
+  const factory = new hre.ethers.ContractFactory(artifact.abi, artifact.bytecode, signer);
+  const contract = await factory.deploy(/* constructor args */);
+  await contract.waitForDeployment();
+  return contract as Contract;
+};
+
+export const getFixedMockContract = async (hre: HardhatRuntimeEnvironment, artifact: MockArtifact) => {
+  if (!artifact.isFixed) {
+    throw new Error('Artifact is not fixed');
+  }
+  return await hre.ethers.getContractAt(artifact.abi, artifact.fixedAddress);
+};
+
+// Testing utils
 
 const mock_checkIsTestnet = async (fnName: string, provider: HardhatEthersProvider | ethers.JsonRpcProvider) => {
   // Testnet is checked by testing if MockZkVerifier is deployed