From ff6782884898682c44eb797dd996fd2a11150313 Mon Sep 17 00:00:00 2001
From: PacificYield <173040337+PacificYield@users.noreply.github.com>
Date: Tue, 12 Nov 2024 12:05:48 +0100
Subject: [PATCH] refactor: Add totalSupply in constructor

---
 contracts/governance/Comp.sol          | 16 +++++++++-------
 contracts/test/governance/TestComp.sol |  5 +++--
 test/governance/Comp.fixture.ts        |  3 ++-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/contracts/governance/Comp.sol b/contracts/governance/Comp.sol
index 0f19439..5c71b8d 100644
--- 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
index 1b5290a..200faf2 100644
--- 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
index 62b0af8..889bb98 100644
--- 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;
 }