refactor: Add totalSupply in constructor

zama-ai · Nov 12, 2024 · ff67828 · ff67828
1 parent 01a8bef
commit ff67828
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/contracts/governance/Comp.sol b/contracts/governance/Comp.sol
@@ -85,19 +85,21 @@ abstract contract Comp is IComp, EncryptedERC20, EIP712, Ownable2Step {
     euint64 private _EUINT64_ZERO;
 
     /**
-     * @param owner_    Owner address.
-     * @param name_     Token name.
-     * @param symbol_   Token symbol.
-     * @param version_  Version (e.g. "0.1", "1.0").
+     * @param owner_        Owner address.
+     * @param name_         Token name.
+     * @param symbol_       Token symbol.
+     * @param version_      Version (e.g. "0.1", "1.0").
+     * @param totalSupply_  Total supply to mint.
      */
     constructor(
         address owner_,
         string memory name_,
         string memory symbol_,
-        string memory version_
+        string memory version_,
+        uint64 totalSupply_
     ) EncryptedERC20(name_, symbol_) EIP712(name_, version_) Ownable(owner_) {
-        _unsafeMint(owner_, TFHE.asEuint64(10000000e6)); /// 10 million Comp
-        _totalSupply = 10000000e6;
+        _unsafeMint(owner_, TFHE.asEuint64(totalSupply_));
+        _totalSupply = totalSupply_;
 
         /// @dev Define the constant in the storage.
         _EUINT64_ZERO = TFHE.asEuint64(0);

diff --git a/contracts/test/governance/TestComp.sol b/contracts/test/governance/TestComp.sol
@@ -9,8 +9,9 @@ contract TestComp is DefaultFHEVMConfig, Comp {
         address owner_,
         string memory name_,
         string memory symbol_,
-        string memory version_
-    ) Comp(owner_, name_, symbol_, version_) {
+        string memory version_,
+        uint64 totalSupply_
+    ) Comp(owner_, name_, symbol_, version_, totalSupply_) {
         //
     }
 }
diff --git a/test/governance/Comp.fixture.ts b/test/governance/Comp.fixture.ts
@@ -1,3 +1,4 @@
+import { parseUnits } from "ethers";
 import { ethers } from "hardhat";
 
 import type { TestComp } from "../../types";
@@ -9,7 +10,7 @@ export async function deployCompFixture(signers: Signers): Promise<TestComp> {
   const contractFactory = await ethers.getContractFactory("TestComp");
   const contract = await contractFactory
     .connect(signers.alice)
-    .deploy(signers.alice.address, "CompoundZama", "COMP", "1.0");
+    .deploy(signers.alice.address, "CompoundZama", "COMP", "1.0", parseUnits("10000000", 6));
   await contract.waitForDeployment();
   return contract;
 }