Enhance VestingManager Functionality and Testing

contracts/XXXToken.sol

@@ -53,7 +53,7 @@ contract XXXToken is Initializable,
      * @dev Initializes the contract replacing the constructor for upgradeable contracts
      * Sets up roles and configures token parameters
      */
-    function initialize() public initializer {
+    function initialize() external initializer {
         // Initialize ERC20 with name and symbol
         __ERC20_init("XXX", "XXX");
 

contracts/XXXTokenVault.sol

@@ -79,7 +79,7 @@ contract TokenVault is Initializable,
      * @dev Initializes the contract replacing the constructor for upgradeable contracts
      * @param _ttnToken Address of the XXXToken contract
      */
-    function initialize(address _ttnToken) public initializer {
+    function initialize(address _ttnToken) external initializer {
          if (_ttnToken == address(0)) revert ZeroAddress("token");
         __AccessControl_init();
         __UUPSUpgradeable_init();

contracts/XXXVestingManager.sol

@@ -112,7 +112,7 @@ contract VestingManager is Initializable,
      * @param _ttnToken Address of the XXXToken contract
      * @param _tokenVault Address of the TokenVault contract
      */
-    function initialize(address _ttnToken, address _tokenVault) public initializer {
+    function initialize(address _ttnToken, address _tokenVault) external initializer {
        if (_ttnToken == address(0)) revert ZeroAddress("token");
        if (_tokenVault == address(0)) revert ZeroAddress("vault");
         __AccessControl_init();
@@ -200,7 +200,7 @@ contract VestingManager is Initializable,
      * @param scheduleId ID of the vesting schedule
      * @return The amount of tokens that can be released
      */
-    function computeReleasableAmount(uint256 scheduleId) public view returns (uint256) {
+    function computeReleasableAmount(uint256 scheduleId) internal view returns (uint256) {
         if (scheduleId == 0 || scheduleId > _vestingScheduleCounter) revert InvalidScheduleId();
 
         VestingSchedule storage schedule = vestingSchedules[scheduleId];

contracts/upgrades/XXXTokenV2.sol

@@ -24,12 +24,22 @@ contract XXXTokenV2 is XXXToken {
     /// @notice The version number of this contract implementation
     uint256 public version;
 
+
+
     /**
-     * @dev Initializes the V2 contract with a version number.
-     * This function is called after the upgrade to set up the new state variables.
-     * The reinitializer modifier ensures this can only be called once after the upgrade.
+     * @dev Initializes V2 functionality
+     * This is called during the upgrade process
+     * @custom:oz-upgrades-validate-as-initializer
      */
-    function initializeV2() public reinitializer(2) {
+    function initializeV2() external reinitializer(2) {
+        // Initialize parent contracts
+        __ERC20_init("XXX Token", "XXX");
+        __ERC20Capped_init(1000000000 * 10**18); // 1 billion tokens
+        __Ownable_init(msg.sender);
+        __AccessControl_init();
+        __Pausable_init();
+        __UUPSUpgradeable_init();
+
         version = 2;
     }
 } 

contracts/upgrades/XXXTokenVaultV2.sol

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: MIT
+/**
+ * @title XXXTokenVaultV2
+ * @dev TESTING PURPOSES ONLY - DO NOT USE IN PRODUCTION
+ * This is a test implementation of TokenVault V2 to demonstrate upgrade functionality.
+ * It adds version tracking and additional allocation tracking for testing purposes.
+ * This contract should not be used in production environments.
+ */
+
+pragma solidity ^0.8.24;
+
+import "../XXXTokenVault.sol";
+import "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+
+contract XXXTokenVaultV2 is TokenVault {
+    // Version tracking for upgrade testing
+    uint256 public version;
+    // Additional counter for V2 allocations
+    uint256 public totalAllocatedV2;
+
+
+
+
+    /**
+     * @dev Initializes V2 functionality
+     * This is called during the upgrade process
+     * @custom:oz-upgrades-validate-as-initializer
+     */
+    function initializeV2() external reinitializer(2) {
+        // Initialize parent contracts
+        __ReentrancyGuard_init();
+        __AccessControl_init();
+        __Pausable_init();
+        __UUPSUpgradeable_init();
+
+
+
+        version = 2;
+    }
+
+    /**
+     * @dev Returns the current version number
+     * @return The version number (2 for V2)
+     */
+    function getVersion() external view returns (uint256) {
+        return version;
+    }
+
+    /**
+     * @dev Returns the total amount allocated through V2
+     * @return The total amount allocated using V2 functions
+     */
+    function getTotalAllocatedV2() external view returns (uint256) {
+        return totalAllocatedV2;
+    }
+
+    /**
+     * @dev Creates a new allocation and tracks it in V2 counter
+     * @param beneficiary Address to receive allocated tokens
+     * @param amount Total amount of tokens to allocate
+     * @return allocationId Unique identifier for the allocation
+     */
+    function createAllocationV2(address beneficiary, uint256 amount) external onlyRole(DEFAULT_ADMIN_ROLE) returns (uint256) {
+        // Input validation
+        require(beneficiary != address(0), "Invalid beneficiary");
+        require(amount > 0, "Amount must be greater than 0");
+
+        // Track V2 allocation
+        totalAllocatedV2 += amount;
+
+        // Call base contract's allocation function
+        return this.createAllocation(beneficiary, amount);
+    }
+} 

contracts/upgrades/XXXVestingManagerV2.sol

@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: MIT
+/**
+ * @title XXXVestingManagerV2
+ * @dev TESTING PURPOSES ONLY - DO NOT USE IN PRODUCTION
+ * This is a test implementation of VestingManager V2 to demonstrate upgrade functionality.
+ * It adds version tracking and additional vesting tracking for testing purposes.
+ * This contract should not be used in production environments.
+ */
+
+pragma solidity ^0.8.24;
+
+import "../XXXVestingManager.sol";
+import "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+
+
+contract XXXVestingManagerV2 is VestingManager {
+    // Version tracking for upgrade testing
+    uint256 public version;
+    // Additional counter for V2 vesting schedules
+    uint256 public totalVestedV2;
+
+    /**
+     * @dev Initializes V2 functionality
+     * This is called during the upgrade process
+     * @custom:oz-upgrades-validate-as-initializer
+     */
+    function initializeV2() external reinitializer(2) {
+        // Initialize parent contracts
+        __ReentrancyGuard_init();
+        __AccessControl_init();
+        __Pausable_init();
+        __UUPSUpgradeable_init();
+
+        // Set version to 2
+        version = 2;
+    }
+
+    /**
+     * @dev Returns the current version number
+     * @return The version number (2 for V2)
+     */
+    function getVersion() external view returns (uint256) {
+        return version;
+    }
+
+    /**
+     * @dev Returns the total amount vested through V2
+     * @return The total amount vested using V2 functions
+     */
+    function getTotalVestedV2() external view returns (uint256) {
+        return totalVestedV2;
+    }
+
+    /**
+     * @dev Creates a new vesting schedule and tracks it in V2 counter
+     * @param beneficiary Address to receive vested tokens
+     * @param amount Total amount of tokens to vest
+     * @param startTime Unix timestamp when vesting begins
+     * @param cliffDuration Duration in seconds until first tokens unlock
+     * @param duration Total duration of vesting in seconds
+     * @param slicePeriodSeconds Duration of each vesting slice in seconds
+     * @return scheduleId Unique identifier for the vesting schedule
+     */
+    function createVestingScheduleV2(
+        address beneficiary,
+        uint256 amount,
+        uint256 startTime,
+        uint256 cliffDuration,
+        uint256 duration,
+        uint256 slicePeriodSeconds
+    ) external onlyRole(VESTING_ADMIN_ROLE) returns (uint256) {
+        // Input validation
+        require(beneficiary != address(0), "Invalid beneficiary");
+        require(amount > 0, "Amount must be greater than 0");
+
+        // Track V2 vesting
+        totalVestedV2 += amount;
+
+        // Call base contract's vesting schedule function
+        return this.createVestingSchedule(
+            beneficiary,
+            amount,
+            startTime,
+            cliffDuration,
+            duration,
+            slicePeriodSeconds
+        );
+    }
+} 

test/TokenVault.test.ts

@@ -76,7 +76,7 @@ describe("TokenVault", function () {
     });
   });
 
-  describe("Vesting Manager", function () {
+  describe("Set Vesting Manager", function () {
     it("Should allow admin to set vesting manager", async function () {
       await vault.setVestingManager(vestingManager.address);
       expect(await vault.vestingManager()).to.equal(vestingManager.address);
@@ -236,4 +236,45 @@ describe("TokenVault", function () {
       ).to.be.revertedWithCustomError(vault, "EnforcedPause");
     });
   });
+
+  describe("Token Minting", function () {
+    it("Should allow minting up to max supply", async function () {
+      const maxSupply = await token.MAX_SUPPLY();
+      const halfSupply = maxSupply / 2n;
+
+      // Mint half the supply
+      await vault.connect(owner).createAllocation(beneficiary1.address, halfSupply);
+      await vault.connect(owner).createAllocation(beneficiary2.address, halfSupply);
+
+      // Verify total supply
+      expect(await token.totalSupply()).to.equal(maxSupply);
+    });
+
+    it("Should not allow minting beyond max supply", async function () {
+      const maxSupply = await token.MAX_SUPPLY();
+      const halfSupply = maxSupply / 2n;
+
+      // Mint half the supply
+      await vault.connect(owner).createAllocation(beneficiary1.address, halfSupply);
+      await vault.connect(owner).createAllocation(beneficiary2.address, halfSupply);
+
+      // Try to mint one more token
+      await expect(
+        vault.connect(owner).createAllocation(beneficiary1.address, 1)
+      ).to.be.revertedWithCustomError(token, "MaxSupplyExceeded");
+    });
+
+    it("Should not allow minting that would exceed max supply", async function () {
+      const maxSupply = await token.MAX_SUPPLY();
+      const halfSupply = maxSupply / 2n;
+
+      // Mint half the supply
+      await vault.connect(owner).createAllocation(beneficiary1.address, halfSupply);
+
+      // Try to mint more than the remaining supply
+      await expect(
+        vault.connect(owner).createAllocation(beneficiary2.address, halfSupply + 1n)
+      ).to.be.revertedWithCustomError(token, "MaxSupplyExceeded");
+    });
+  });
 });