add upgradable test to the token

Author: KoxyG

contracts/XXXTokenV2.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import "./XXXToken.sol";
+
+/**
+ * @title XXXTokenV2
+ * @dev This is a demonstration contract for testing the upgradeability of XXXToken.
+ * It serves as a simple example of how to implement contract upgrades using the UUPS pattern.
+ * 
+ * Key features:
+ * - Inherits all functionality from XXXToken
+ * - Adds a version number to demonstrate upgrade success
+ * - Uses reinitializer to safely initialize new state variables
+ * 
+ * This contract is primarily used for testing purposes to verify that:
+ * 1. The upgrade process works correctly
+ * 2. All existing functionality is preserved
+ * 3. New functionality can be added safely
+ * 
+ * @custom:oz-upgrades-validate-as-initializer
+ */
+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.
+     */
+    function initializeV2() public reinitializer(2) {
+        version = 2;
+    }
+} 

test/TokenVault.test.ts

@@ -41,11 +41,13 @@ describe("TokenVault", function () {
     const ALLOCATOR_ROLE = ethers.keccak256(ethers.toUtf8Bytes("ALLOCATOR_ROLE"));
     const AIRDROP_ROLE = ethers.keccak256(ethers.toUtf8Bytes("AIRDROP_ROLE"));
     const UPGRADER_ROLE = ethers.keccak256(ethers.toUtf8Bytes("UPGRADER_ROLE"));
+    const DEFAULT_ADMIN_ROLE = "0x0000000000000000000000000000000000000000000000000000000000000000";
 
     // Grant MINTER_ROLE to the vault contract
     await token.grantRole(MINTER_ROLE, await vault.getAddress());
 
     // Grant all other roles to the owner
+    await vault.grantRole(DEFAULT_ADMIN_ROLE, owner.address);
     await vault.grantRole(ALLOCATOR_ROLE, owner.address);
     await vault.grantRole(AIRDROP_ROLE, owner.address);
     await vault.grantRole(UPGRADER_ROLE, owner.address);

test/XXXToken.test.ts

@@ -5,7 +5,7 @@ import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
 import { Contract } from "ethers";
 
 // Import the contract type
-import { XXXToken } from "../typechain-types";
+import { XXXToken, XXXTokenV2 } from "../typechain-types";
 
 describe("XXXToken", function () {
   // Test roles
@@ -157,19 +157,29 @@ describe("XXXToken", function () {
     });
   });
 
-  // describe("Upgrading", function () {
-  //   it("Should allow upgrader to upgrade the contract", async function () {
-  //     const XXXTokenV2 = await ethers.getContractFactory("XXXTokenV2");
-  //     await expect(
-  //       token.connect(upgrader).upgradeTo(await XXXTokenV2.getAddress())
-  //     ).to.not.be.reverted;
-  //   });
-
-  //   it("Should not allow non-upgrader to upgrade the contract", async function () {
-  //     const XXXTokenV2 = await ethers.getContractFactory("XXXTokenV2");
-  //     await expect(
-  //       token.connect(user1).upgradeTo(await XXXTokenV2.getAddress())
-  //     ).to.be.revertedWithCustomError(token, "AccessControlUnauthorizedAccount");
-  //   });
-  // });
+  describe("Upgrading", function () {
+    it("Should allow upgrader to upgrade the contract", async function () {
+      const XXXTokenV2Factory = await ethers.getContractFactory("XXXTokenV2");
+      const upgraded = await upgrades.upgradeProxy(await token.getAddress(), XXXTokenV2Factory) as unknown as XXXTokenV2;
+      await upgraded.initializeV2();
+      expect(await upgraded.version()).to.equal(2);
+    });
+
+    it("Should not allow initializing V2 twice", async function () {
+      const XXXTokenV2Factory = await ethers.getContractFactory("XXXTokenV2");
+      const upgraded = await upgrades.upgradeProxy(await token.getAddress(), XXXTokenV2Factory) as unknown as XXXTokenV2;
+      await upgraded.initializeV2();
+      
+      await expect(
+        upgraded.initializeV2()
+      ).to.be.revertedWithCustomError(upgraded, "InvalidInitialization");
+    });
+
+    it("Should not allow non-upgrader to upgrade the contract", async function () {
+      const XXXTokenV2Factory = await ethers.getContractFactory("XXXTokenV2");
+      await expect(
+        upgrades.upgradeProxy(await token.getAddress(), XXXTokenV2Factory.connect(user1))
+      ).to.be.revertedWithCustomError(token, "AccessControlUnauthorizedAccount");
+    });
+  });
 });