@cofhe/mock-contracts 0.2.0 → 0.3.0

package/CHANGELOG.md

@@ -1,5 +1,31 @@
 # @cofhe/mock-contracts Changelog
 
+## 0.3.0
+
+### Minor Changes
+
+- 35024b6: Remove `sdk` from function names and exported types. Rename:
+
+  - `createCofhesdkConfig` -> `createCofheConfig`
+  - `createCofhesdkClient` -> `createCofheClient`
+  - `hre.cofhesdk.*` -> `hre.cofhe.*`
+  - `hre.cofhesdk.createCofheConfig()` → `hre.cofhe.createConfig()`
+  - `hre.cofhesdk.createCofheClient()` → `hre.cofhe.createClient()`
+  - `hre.cofhesdk.createBatteriesIncludedCofheClient()` → `hre.cofhe.createClientWithBatteries()`
+
+- 29c2401: implement decrypt-with-proof flows and related tests:
+
+  - Implement production `decryptForTx` backed by Threshold Network `POST /decrypt`, with explicit permit vs global-allowance selection.
+  - Rename mocks “Query Decrypter” -> “Threshold Network” and update SDK constants/contracts/artifacts accordingly.
+  - Extend mock contracts + hardhat plugin to publish & verify decryption results on-chain, and add end-to-end integration tests.
+
+## 0.2.1
+
+### Patch Changes
+
+- ac47e2f: Add `PermitUtils.checkValidityOnChain` to validate permits against the on-chain deployed ACL (source of truth).
+- 0000d5e: Mock contracts deployed to alternate fixed addresses to avoid collision with hardhat pre-compiles.
+
 ## 0.2.0
 
 ### Minor Changes

package/README.md

@@ -1,4 +1,4 @@
-# @fhenixprotocol/cofhe-mock-contracts [![NPM Package][npm-badge]][npm] [![License: MIT][license-badge]][license]
+# cofhe/mock-contracts [![NPM Package][npm-badge]][npm] [![License: MIT][license-badge]][license]
 
 [npm]: https://www.npmjs.com/package/@fhenixprotocol/cofhe-mock-contracts
 [npm-badge]: https://img.shields.io/npm/v/@fhenixprotocol/cofhe-mock-contracts.svg
@@ -11,7 +11,7 @@ A mock smart contract library for testing CoFHE (Confidential Computing Framewor
 
 - Mock implementations of core CoFHE contracts:
   - MockTaskManager
-  - MockQueryDecrypter
+  - MockThresholdNetwork
   - MockZkVerifier
   - ACL (Access Control List)
 - Synchronous operation simulation with mock delays
@@ -34,11 +34,67 @@ forge install fhenixprotocol/cofhe-mock-contracts
 
 ## Usages and Integrations
 
-`@cofhe/sdk` is designed to work with mock contracts in a testing / hardhat environment. `cofhesdk/hardhat-plugin` deploys the mock contracts in this repo, and the `cofhesdkClient` detects a testnet chain and interacts correctly using the mocks rather than the true CoFHE coprocessor.
+### Who is this for?
+
+This package is intended for **developers building and testing CoFHE-enabled applications and smart contracts**.
+
+Use these mocks when you want to:
+
+- Run **local tests** without depending on the real CoFHE coprocessor infrastructure.
+- Debug flows end-to-end (encrypt → submit → operate → decrypt) with fast iteration.
+- Assert on results deterministically in CI.
+
+Do **not** use these mocks for production deployments: they intentionally make testing convenient (e.g. storing plaintext on-chain for inspection) and therefore **do not provide real confidentiality guarantees**.
+
+### Hardhat integration vs Foundry integration
+
+Both integrations use the same underlying mock contracts, but they differ in **how mocks get deployed** and **how you interact with them**.
+
+#### Hardhat (recommended for TS/SDK + Solidity tests)
+
+Use this when you are already using **Hardhat** and/or want to run the **TypeScript SDK (`@cofhe/sdk`)** against a local chain.
+
+- The `cofhesdk/hardhat-plugin` watches Hardhat `node` and `test` tasks.
+- It automatically deploys the mocks to the Hardhat network at fixed addresses.
+- The `cofheClient` (created with `createCofheClient(...)`) detects the mocks and routes CoFHE actions to them.
+
+Minimal setup:
+
+```ts
+// hardhat.config.ts
+import 'cofhe-hardhat-plugin';
+
+export default {
+  cofhe: {
+    logMocks: true, // optional
+  },
+};
+```
+
+Run:
+
+```bash
+npx hardhat test
+# or
+npx hardhat node
+```
+
+If you want to assert on plaintext values in Hardhat tests, the plugin exposes helpers like `mock_expectPlaintext(...)` (see the hardhat-plugin README).
+
+#### Foundry (recommended for Solidity-only tests)
+
+Use this when you are writing tests in **Solidity** and running them with `forge test`.
+
+- You typically inherit from the abstract `CoFheTest` helper to deploy/setup the necessary FHE mock environment.
+- You use helper methods to create encrypted inputs and assert their underlying values.
+
+> **Important**: You must set `isolate = true` in your `foundry.toml`. Without this setting, some variables may be used without proper permission checks, which will cause failures on production chains.
+
+`@cofhe/sdk` is designed to work with mock contracts in a testing / hardhat environment. `cofhesdk/hardhat-plugin` deploys the mock contracts in this repo, and the `cofheClient` detects a testnet chain and interacts correctly using the mocks rather than the true CoFHE coprocessor.
 
 When installed and imported in the `hardhat.config.ts`, `cofhesdk/hardhat-plugin` will watch for Hardhat `node` and `test` tasks, and will deploy the mocks to the hardhat testnet chain at fixed addresses.
 
-Once deployed, interaction with the mock contracts is handled by the `cofhesdkClient` (created with `createCofhesdkClient(...)`). The client checks for the existence of mock contracts at known addresses, and if they exist, marks the current connection as a testnet.
+Once deployed, interaction with the mock contracts is handled by the `cofheClient` (created with `createCofheClient(...)`). The client checks for the existence of mock contracts at known addresses, and if they exist, marks the current connection as a testnet.
 
 ## Logging
 
@@ -71,20 +127,18 @@ The mocks are then responsible for mocking two actions:
 1. Creating the signature.
 2. Storing the plaintext value on-chain.
 
-The `MockZkVerifier` contract handles the on-chain storage of encrypted inputs. The signature creation is handled automatically within `cofhesdkClient.encryptInputs` when executing against a testnet.
+The `MockZkVerifier` contract handles the on-chain storage of encrypted inputs. The signature creation is handled automatically within `cofheClient.encryptInputs` when executing against a testnet.
 
 ### Off-chain Decryption / Sealing
 
-Off-chain decryption is performed by calling the `cofhesdkClient.decryptHandle` function with a valid `ctHash` and a valid `permit` [todo link].
+Off-chain decryption is performed by calling the `cofheClient.decryptHandle` function with a valid `ctHash` and a valid `permit` [todo link].
 
 When interacting with CoFHE this request is routed to the Threshold Network, which will perform the decryption operation, ultimately returning a decrypted result.
 
-When working with the mocks, the `cofhesdkClient` will instead query the `MockQueryDecrypter` contract, which will verify the request `permit`, and return the decrypted result.
+When working with the mocks, the `cofheClient` will instead query the `MockThresholdNetwork` contract, which will verify the request `permit`, and return the decrypted result.
 
 ### Using Foundry
 
-> **Important**: You must set `isolate = true` in your `foundry.toml`. Without this setting, some variables may be used without proper permission checks, which will cause failures on production chains.
-
 Use abstract CoFheTest contract to automatically deploy all necessary FHE contracts for testing.
 
 CoFheTest also exposes useful test methods such as

package/contracts/MockACL.sol

@@ -335,4 +335,8 @@ contract MockACL is MockPermissioned {
   ) public view withPermission(permission) returns (bool) {
     return isAllowed(handle, permission.issuer);
   }
+
+  function checkPermitValidity(Permission memory permission) public view withPermission(permission) returns (bool) {
+    return true;
+  }
 }

package/contracts/MockTaskManager.sol

@@ -30,6 +30,7 @@ error OnlyAggregatorAllowed(address caller);
 
 // Operation-specific errors
 error RandomFunctionNotSupported();
+error NotImplemented();
 
 library TMCommon {
   uint256 private constant HASH_MASK_FOR_METADATA = type(uint256).max - type(uint16).max; // 2 bytes reserved for metadata
@@ -150,6 +151,7 @@ library TMCommon {
   function getSecAndTypeFromHash(uint256 hash) internal pure returns (uint256) {
     return uint256((SHIFTED_TYPE_MASK | SECURITY_ZONE_MASK) & hash);
   }
+
   function isTriviallyEncryptedFromHash(uint256 hash) internal pure returns (bool) {
     return (hash & TRIVIALLY_ENCRYPTED_MASK) == TRIVIALLY_ENCRYPTED_MASK;
   }
@@ -218,6 +220,9 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
 
   address private verifierSigner;
 
+  // Signer address for decrypt result verification (threshold network's signing key)
+  address public decryptResultSigner;
+
   // Storage contract for plaintext results of decrypt operations
   // PlaintextsStorage public plaintextsStorage;
 
@@ -288,7 +293,7 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     // NOTE: MOCK
     if (block.timestamp < _decryptResultReadyTimestamp[ctHash]) return (0, false);
 
-    return (_get(ctHash), true);
+    return (_decryptResult[ctHash], true);
   }
 
   function checkAllowed(uint256 ctHash) internal view {
@@ -460,6 +465,50 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     emit ProtocolNotification(ctHash, operation, errorMessage);
   }
 
+  /// @notice Publish a signed decrypt result to the chain
+  /// @dev Anyone with a valid signature from the decrypt network can call this
+  function publishDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) external {
+    _verifyDecryptResult(ctHash, result, signature);
+    _decryptResultReady[ctHash] = true;
+    _decryptResult[ctHash] = result;
+
+    _decryptResultReadyTimestamp[ctHash] = uint64(block.timestamp);
+    emit DecryptionResult(ctHash, result, msg.sender);
+  }
+
+  function publishDecryptResultBatch(
+    uint256[] calldata ctHashes,
+    uint256[] calldata results,
+    bytes[] calldata signatures
+  ) external {
+    // Mock implementation: process each result individually
+    // Note: publishDecryptResult is defined later in the contract
+    for (uint256 i = 0; i < ctHashes.length; i++) {
+      // Call via external interface to avoid forward reference issue
+      this.publishDecryptResult(ctHashes[i], results[i], signatures[i]);
+    }
+  }
+
+  function _verifyDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) private view {
+    if (decryptResultSigner == address(0)) revert InvalidAddress();
+    bytes32 messageHash = _computeDecryptResultHash(ctHash, result);
+    (address recovered, ECDSA.RecoverError err, ) = ECDSA.tryRecover(messageHash, signature);
+
+    if (err != ECDSA.RecoverError.NoError || recovered == address(0)) {
+      revert InvalidSignature();
+    }
+
+    if (recovered != decryptResultSigner) {
+      revert InvalidSigner(recovered, decryptResultSigner);
+    }
+  }
+
+  /// @notice Format: result (32) || enc_type (4) || chain_id (8) || ct_hash (32) = 76 bytes
+  function _computeDecryptResultHash(uint256 ctHash, uint256 result) private view returns (bytes32) {
+    uint8 encryptionType = TMCommon.getUintTypeFromHash(ctHash);
+    return keccak256(abi.encodePacked(result, uint32(encryptionType), uint64(block.chainid), bytes32(ctHash)));
+  }
+
   function verifyType(uint8 ctType, uint8 desiredType) internal pure {
     if (ctType != desiredType) {
       revert InvalidInputType(ctType, desiredType);
@@ -550,6 +599,12 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
     verifierSigner = signer;
   }
 
+  /// @notice Set the authorized signer for decrypt results
+  /// @param signer The new signer address (must be non-zero for publishDecryptResult to work)
+  function setDecryptResultSigner(address signer) external onlyOwner {
+    decryptResultSigner = signer;
+  }
+
   function setSecurityZoneMax(int32 securityZone) external onlyOwner {
     if (securityZone < securityZoneMin) {
       revert InvalidSecurityZone(securityZone, securityZoneMin, securityZoneMax);
@@ -590,4 +645,34 @@ contract MockTaskManager is ITaskManager, MockCoFHE {
   function isAllowedWithPermission(Permission memory permission, uint256 handle) public view returns (bool) {
     return acl.isAllowedWithPermission(permission, handle);
   }
+
+  // Stub implementations for new ITaskManager interface methods (inc PR #48)
+
+  function createRandomTask(uint8 returnType, uint256 seed, int32 securityZone) external returns (uint256) {
+    // Mock implementation: just return a pseudo-random hash based on seed
+    return uint256(keccak256(abi.encode(returnType, seed, securityZone, block.timestamp)));
+  }
+
+  function isPubliclyAllowed(uint256 ctHash) external view returns (bool) {
+    revert NotImplemented();
+  }
+
+  function verifyDecryptResult(uint256 ctHash, uint256 result, bytes calldata signature) external view returns (bool) {
+    // Mock implementation: verify signature using the verifier signer
+    bytes32 digest = keccak256(abi.encodePacked(result));
+    return ECDSA.recover(digest, signature) == verifierSigner;
+  }
+
+  function verifyDecryptResultSafe(
+    uint256 ctHash,
+    uint256 result,
+    bytes calldata signature
+  ) external view returns (bool) {
+    // Same as verifyDecryptResult for mock
+    try this.verifyDecryptResult(ctHash, result, signature) returns (bool valid) {
+      return valid;
+    } catch {
+      return false;
+    }
+  }
 }

package/contracts/{MockQueryDecrypter.sol → MockThresholdNetwork.sol}

@@ -1,13 +1,17 @@
 // SPDX-License-Identifier: BSD-3-Clause-Clear
 // solhint-disable one-contract-per-file
 
+// NOTE: This file was renamed from `MockQueryDecrypter.sol` to `MockThresholdNetwork.sol`
+// to better reflect that it mocks threshold-network-style decryption/sealing behavior
+// in local/testing environments.
+
 pragma solidity >=0.8.19 <0.9.0;
 
 import { MockACL } from './MockACL.sol';
 import { MockTaskManager } from './MockTaskManager.sol';
 import { Permission, MockPermissioned } from './Permissioned.sol';
 
-contract MockQueryDecrypter {
+contract MockThresholdNetwork {
   MockTaskManager public mockTaskManager;
   MockACL public mockAcl;
 
@@ -111,6 +115,68 @@ contract MockQueryDecrypter {
     return (true, '', seal(value, permission.sealingKey));
   }
 
+  // DECRYPT FOR TX
+
+  function _isAllowedWithPermit(
+    uint256 ctHash,
+    Permission memory permission
+  ) internal view returns (bool isAllowed, string memory error) {
+    try mockTaskManager.isAllowedWithPermission(permission, ctHash) returns (bool _isAllowed) {
+      isAllowed = _isAllowed;
+    } catch Error(string memory reason) {
+      return (false, reason);
+    } catch Panic(uint /*errorCode*/) {
+      return (false, 'Panic');
+    } catch (bytes memory lowLevelData) {
+      return (false, decodeLowLevelReversion(lowLevelData));
+    }
+
+    if (!isAllowed) return (false, 'NotAllowed');
+    return (true, '');
+  }
+
+  function _isAllowedWithoutPermit(uint256 ctHash) internal view returns (bool isAllowed, string memory error) {
+    try mockAcl.globalAllowed(ctHash) returns (bool _isAllowed) {
+      isAllowed = _isAllowed;
+    } catch Error(string memory reason) {
+      return (false, reason);
+    } catch Panic(uint /*errorCode*/) {
+      return (false, 'Panic');
+    } catch (bytes memory lowLevelData) {
+      return (false, decodeLowLevelReversion(lowLevelData));
+    }
+
+    if (!isAllowed) return (false, 'NotAllowed');
+    return (true, '');
+  }
+
+  /// @notice Decrypt a ciphertext for a transaction using a permit.
+  function decryptForTxWithPermit(
+    uint256 ctHash,
+    Permission memory permission
+  ) public view returns (bool allowed, string memory error, uint256 decryptedValue) {
+    if (permission.issuer == address(0)) {
+      return (false, 'PermissionMissing', 0);
+    }
+
+    (bool isAllowed, string memory err) = _isAllowedWithPermit(ctHash, permission);
+    if (!isAllowed) return (false, err, 0);
+
+    uint256 value = mockTaskManager.mockStorage(ctHash);
+    return (true, '', value);
+  }
+
+  /// @notice Decrypt a ciphertext for a transaction using global allowance (no permit).
+  function decryptForTxWithoutPermit(
+    uint256 ctHash
+  ) public view returns (bool allowed, string memory error, uint256 decryptedValue) {
+    (bool isAllowed, string memory err) = _isAllowedWithoutPermit(ctHash);
+    if (!isAllowed) return (false, err, 0);
+
+    uint256 value = mockTaskManager.mockStorage(ctHash);
+    return (true, '', value);
+  }
+
   // UTIL
 
   function decodeLowLevelReversion(bytes memory data) public pure returns (string memory error) {

package/contracts/TestBed.sol

@@ -3,14 +3,31 @@ pragma solidity ^0.8.13;
 
 import '@fhenixprotocol/cofhe-contracts/FHE.sol';
 
+/// @title TestBed
+/// @notice Minimal contract used to smoke-test CoFHE/FHE flows (mocks + permissions + decrypt plumbing).
+/// @dev This contract is intentionally simple and is primarily a convenience for Hardhat-based local
+///      development: the Hardhat plugin can deploy it automatically when `deployTestBed` is enabled.
+///
+///      In Foundry flows, nothing deploys or depends on this contract automatically — tests must
+///      instantiate it explicitly (this repository includes such Foundry tests). You may deploy/use it
+///      in `forge test` if you want a known-good reference target.
+///
+///      Important concepts:
+///      - `eNumber` is an on-chain encrypted handle type (`euint32`). In real CoFHE, this represents a ciphertext.
+///      - `numberHash` stores the unwrapped handle (`ctHash`) as a uint256 for easy inspection/assertions.
+///      - After every state update we call `FHE.allowThis(...)` and `FHE.allowSender(...)` so the contract
+///        and the transaction sender can continue operating on / decrypting the updated handle.
 contract TestBed {
   euint32 public eNumber;
-  uint256 public numberHash;
+  bytes32 public numberHash;
 
+  /// @notice Marker used by deploy scripts/tests to confirm the contract is deployed.
   function exists() public pure returns (bool) {
     return true;
   }
 
+  /// @notice Sets `eNumber` from an encrypted input struct.
+  /// @dev Typically used when testing client-side encryption flows.
   function setNumber(InEuint32 memory inNumber) public {
     eNumber = FHE.asEuint32(inNumber);
     numberHash = euint32.unwrap(eNumber);
@@ -18,19 +35,23 @@ contract TestBed {
     FHE.allowSender(eNumber);
   }
 
-  function setNumberTrivial(uint256 inNumber) public {
+  /// @notice Convenience setter that casts a plaintext value into an encrypted handle.
+  /// @dev Useful for quick smoke tests that don't need pre-encryption.
+  function setNumberTrivial(uint32 inNumber) public {
     eNumber = FHE.asEuint32(inNumber);
     numberHash = euint32.unwrap(eNumber);
     FHE.allowThis(eNumber);
     FHE.allowSender(eNumber);
   }
 
+  /// @notice Increments `eNumber` by 1 using FHE arithmetic.
   function increment() public {
     eNumber = FHE.add(eNumber, FHE.asEuint32(1));
     FHE.allowThis(eNumber);
     FHE.allowSender(eNumber);
   }
 
+  /// @notice Adds an encrypted input to `eNumber`.
   function add(InEuint32 memory inNumber) public {
     eNumber = FHE.add(eNumber, FHE.asEuint32(inNumber));
 
@@ -38,6 +59,7 @@ contract TestBed {
     FHE.allowSender(eNumber);
   }
 
+  /// @notice Subtracts an encrypted input from `eNumber`, clamped to 0 to avoid underflow.
   function sub(InEuint32 memory inNumber) public {
     euint32 inAsEuint32 = FHE.asEuint32(inNumber);
     euint32 eSubOrZero = FHE.select(FHE.lte(inAsEuint32, eNumber), inAsEuint32, FHE.asEuint32(0));
@@ -46,21 +68,31 @@ contract TestBed {
     FHE.allowSender(eNumber);
   }
 
+  /// @notice Multiplies `eNumber` by an encrypted input.
   function mul(InEuint32 memory inNumber) public {
     eNumber = FHE.mul(eNumber, FHE.asEuint32(inNumber));
     FHE.allowThis(eNumber);
     FHE.allowSender(eNumber);
   }
 
+  /// @notice Requests decryption of `eNumber`.
+  /// @dev In real CoFHE this is asynchronous; in mocks it is simulated.
   function decrypt() public {
     FHE.decrypt(eNumber);
   }
 
+  /// @notice Reads a decryption result (reverts if not ready depending on implementation).
   function getDecryptResult(euint32 input1) public view returns (uint32) {
     return FHE.getDecryptResult(input1);
   }
 
+  /// @notice Reads a decryption result safely, returning a readiness flag.
   function getDecryptResultSafe(euint32 input1) public view returns (uint32 value, bool decrypted) {
     return FHE.getDecryptResultSafe(input1);
   }
+
+  /// @notice Publishes a decrypt result for an encrypted handle.
+  function publishDecryptResult(euint32 input, uint32 result, bytes memory signature) external {
+    FHE.publishDecryptResult(input, result, signature);
+  }
 }

package/contracts/foundry/CoFheTest.sol

@@ -9,7 +9,7 @@ import { MockZkVerifier } from '../MockZkVerifier.sol';
 import { MockZkVerifierSigner } from './MockZkVerifierSigner.sol';
 import { MessageHashUtils } from '@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol';
 import { Permission, PermissionUtils } from '../Permissioned.sol';
-import { MockQueryDecrypter } from '../MockQueryDecrypter.sol';
+import { MockThresholdNetwork } from '../MockThresholdNetwork.sol';
 import { SIGNER_ADDRESS } from '../MockCoFHE.sol';
 
 abstract contract CoFheTest is Test {
@@ -17,11 +17,13 @@ abstract contract CoFheTest is Test {
   MockZkVerifier public mockZkVerifier;
   MockZkVerifierSigner public mockZkVerifierSigner;
   MockACL public mockAcl;
-  MockQueryDecrypter public mockQueryDecrypter;
+  MockThresholdNetwork public mockThresholdNetwork;
 
-  address constant QUERY_DECRYPTER_ADDRESS = address(512);
-  address constant ZK_VERIFIER_ADDRESS = address(256);
-  address constant ZK_VERIFIER_SIGNER_ADDRESS = address(257);
+  // Keep these in sync with `packages/sdk/core/consts.ts`
+  address constant ZK_VERIFIER_ADDRESS = 0x0000000000000000000000000000000000005001;
+  address constant THRESHOLD_NETWORK_ADDRESS = 0x0000000000000000000000000000000000005002;
+  // SDK exposes this as `MOCKS_ZK_VERIFIER_SIGNER_ADDRESS`
+  address constant ZK_VERIFIER_SIGNER_ADDRESS = SIGNER_ADDRESS;
   address public constant ACL_ADDRESS = 0xa6Ea4b5291d044D93b73b3CFf3109A1128663E8B;
 
   bool private _log = false;
@@ -75,12 +77,12 @@ abstract contract CoFheTest is Test {
     mockZkVerifierSigner = MockZkVerifierSigner(ZK_VERIFIER_SIGNER_ADDRESS);
     vm.label(address(mockZkVerifierSigner), 'MockZkVerifierSigner');
 
-    // QUERY DECRYPTER
+    // THRESHOLD NETWORK
 
-    deployCodeTo('MockQueryDecrypter.sol:MockQueryDecrypter', QUERY_DECRYPTER_ADDRESS);
-    mockQueryDecrypter = MockQueryDecrypter(QUERY_DECRYPTER_ADDRESS);
-    vm.label(address(mockQueryDecrypter), 'MockQueryDecrypter');
-    mockQueryDecrypter.initialize(TASK_MANAGER_ADDRESS, address(mockAcl));
+    deployCodeTo('MockThresholdNetwork.sol:MockThresholdNetwork', THRESHOLD_NETWORK_ADDRESS);
+    mockThresholdNetwork = MockThresholdNetwork(THRESHOLD_NETWORK_ADDRESS);
+    vm.label(address(mockThresholdNetwork), 'MockThresholdNetwork');
+    mockThresholdNetwork.initialize(TASK_MANAGER_ADDRESS, address(mockAcl));
 
     // SET LOG OPS
 
@@ -120,31 +122,46 @@ abstract contract CoFheTest is Test {
     assertEq(mockTaskManager.inMockStorage(ctHash), true);
     assertEq(mockTaskManager.mockStorage(ctHash), value);
   }
+
+  function assertHashValue(bytes32 ctHash, uint256 value) public view {
+    assertHashValue(uint256(ctHash), value);
+  }
+
   function assertHashValue(uint256 ctHash, uint256 value, string memory message) public view {
     assertEq(mockTaskManager.inMockStorage(ctHash), true, message);
     assertEq(mockTaskManager.mockStorage(ctHash), value, message);
   }
 
+  function assertHashValue(bytes32 ctHash, uint256 value, string memory message) public view {
+    assertHashValue(uint256(ctHash), value, message);
+  }
+
   // Encrypted types (no message)
 
   function assertHashValue(ebool eValue, bool value) public view {
     assertHashValue(ebool.unwrap(eValue), value ? 1 : 0);
   }
+
   function assertHashValue(euint8 eValue, uint8 value) public view {
     assertHashValue(euint8.unwrap(eValue), value);
   }
+
   function assertHashValue(euint16 eValue, uint16 value) public view {
     assertHashValue(euint16.unwrap(eValue), value);
   }
+
   function assertHashValue(euint32 eValue, uint32 value) public view {
     assertHashValue(euint32.unwrap(eValue), value);
   }
+
   function assertHashValue(euint64 eValue, uint64 value) public view {
     assertHashValue(euint64.unwrap(eValue), value);
   }
+
   function assertHashValue(euint128 eValue, uint128 value) public view {
     assertHashValue(euint128.unwrap(eValue), value);
   }
+
   function assertHashValue(eaddress eValue, address value) public view {
     assertHashValue(eaddress.unwrap(eValue), uint256(uint160(value)));
   }
@@ -154,21 +171,27 @@ abstract contract CoFheTest is Test {
   function assertHashValue(ebool eValue, bool value, string memory message) public view {
     assertHashValue(ebool.unwrap(eValue), value ? 1 : 0, message);
   }
+
   function assertHashValue(euint8 eValue, uint8 value, string memory message) public view {
     assertHashValue(euint8.unwrap(eValue), value, message);
   }
+
   function assertHashValue(euint16 eValue, uint16 value, string memory message) public view {
     assertHashValue(euint16.unwrap(eValue), value, message);
   }
+
   function assertHashValue(euint32 eValue, uint32 value, string memory message) public view {
     assertHashValue(euint32.unwrap(eValue), value, message);
   }
+
   function assertHashValue(euint64 eValue, uint64 value, string memory message) public view {
     assertHashValue(euint64.unwrap(eValue), value, message);
   }
+
   function assertHashValue(euint128 eValue, uint128 value, string memory message) public view {
     assertHashValue(euint128.unwrap(eValue), value, message);
   }
+
   function assertHashValue(eaddress eValue, address value, string memory message) public view {
     assertHashValue(eaddress.unwrap(eValue), uint256(uint160(value)), message);
   }
@@ -258,16 +281,6 @@ abstract contract CoFheTest is Test {
     return abi.decode(abi.encode(createEncryptedInput(Utils.EUINT128_TFHE, value, securityZone, sender)), (InEuint128));
   }
 
-  /**
-   * @notice              Creates an InEuint256 to be used as FHE input. Value is stored in MockCoFHE contract, hash is a pointer to that value.
-   * @param value         Value to encrypt.
-   * @param securityZone  Security zone of the encrypted value.
-   * @return              InEuint256.
-   */
-  function createInEuint256(uint256 value, uint8 securityZone, address sender) public returns (InEuint256 memory) {
-    return abi.decode(abi.encode(createEncryptedInput(Utils.EUINT256_TFHE, value, securityZone, sender)), (InEuint256));
-  }
-
   /**
    * @notice              Creates an InEaddress to be used as FHE input. Value is stored in MockCoFHE contract, hash is a pointer to that value.
    * @param value         Value to encrypt.
@@ -308,10 +321,6 @@ abstract contract CoFheTest is Test {
     return createInEuint128(value, 0, sender);
   }
 
-  function createInEuint256(uint256 value, address sender) public returns (InEuint256 memory) {
-    return createInEuint256(value, 0, sender);
-  }
-
   function createInEaddress(address value, address sender) public returns (InEaddress memory) {
     return createInEaddress(value, 0, sender);
   }
@@ -415,7 +424,7 @@ abstract contract CoFheTest is Test {
     uint256 hostChainId,
     Permission memory permission
   ) public view returns (bool, string memory error, uint256) {
-    return mockQueryDecrypter.queryDecrypt(ctHash, hostChainId, permission);
+    return mockThresholdNetwork.queryDecrypt(ctHash, hostChainId, permission);
   }
 
   function querySealOutput(
@@ -423,10 +432,21 @@ abstract contract CoFheTest is Test {
     uint256 hostChainId,
     Permission memory permission
   ) public view returns (bool, string memory error, bytes32) {
-    return mockQueryDecrypter.querySealOutput(ctHash, hostChainId, permission);
+    return mockThresholdNetwork.querySealOutput(ctHash, hostChainId, permission);
+  }
+
+  function decryptForTxWithoutPermit(uint256 ctHash) public view returns (bool, string memory error, uint256) {
+    return mockThresholdNetwork.decryptForTxWithoutPermit(ctHash);
+  }
+
+  function decryptForTxWithPermit(
+    uint256 ctHash,
+    Permission memory permission
+  ) public view returns (bool, string memory error, uint256) {
+    return mockThresholdNetwork.decryptForTxWithPermit(ctHash, permission);
   }
 
   function unseal(bytes32 hashed, bytes32 key) public view returns (uint256) {
-    return mockQueryDecrypter.unseal(hashed, key);
+    return mockThresholdNetwork.unseal(hashed, key);
   }
 }