Implement mainnet forking support for tests

contracts/MaticX.sol

@@ -77,6 +77,17 @@ contract MaticX is
 		);
 	}
 
+	/**
+	 * @dev Initializes the POL related flow.
+	*/
+	function initializePOL() external onlyRole(DEFAULT_ADMIN_ROLE) {
+		require(polToken != address(0), "Zero POL token address");
+		IERC20Upgradeable(polToken).safeApprove(
+			stakeManager,
+			type(uint256).max
+		);
+	}
+
 	/**
 	 * @dev Sets the BOT's admin role.
 	 * @notice Callable by the admin only.

contracts/interfaces/IPolygonMigration.sol

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.7;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+interface IPolygonMigration {
+	event Migrated(address indexed account, uint256 amount);
+
+	event Unmigrated(
+		address indexed account,
+		address indexed recipient,
+		uint256 amount
+	);
+
+	event UnmigrationLockUpdated(bool lock);
+
+	error UnmigrationLocked();
+
+	error InvalidAddressOrAlreadySet();
+
+	error InvalidAddress();
+
+	function migrate(uint256 amount) external;
+
+	function unmigrate(uint256 amount) external;
+
+	function unmigrateTo(address recipient, uint256 amount) external;
+
+	function unmigrateWithPermit(
+		uint256 amount,
+		uint256 deadline,
+		uint8 v,
+		bytes32 r,
+		bytes32 s
+	) external;
+
+	function updateUnmigrationLock(bool unmigrationLocked) external;
+
+	function burn(uint256 amount) external;
+
+	function matic() external view returns (IERC20 maticToken);
+
+	function polygon() external view returns (IERC20 polygonEcosystemToken);
+
+	function unmigrationLocked()
+		external
+		view
+		returns (bool isUnmigrationLocked);
+
+	function version() external pure returns (string memory version);
+}

hardhat.config.ts

@@ -170,6 +170,11 @@ const config: HardhatUserConfig = {
 			initialBaseFeePerGas: 0, // See https://github.com/sc-forks/solidity-coverage/issues/652#issuecomment-896330136
 			blockGasLimit: 30_000_000,
 			mining,
+			forking: {
+				url: envVars.ROOT_CHAIN_RPC,
+				blockNumber: 20_700_204,
+				enabled: true,
+			},
 		},
 		localhost: {
 			url: "http://127.0.0.1:8545",

package.json

@@ -40,6 +40,7 @@
 	},
 	"type": "commonjs",
 	"devDependencies": {
+		"@nomicfoundation/hardhat-network-helpers": "^1.0.11",
 		"@nomiclabs/hardhat-ethers": "^2.0.2",
 		"@nomiclabs/hardhat-etherscan": "^3.1.8",
 		"@nomiclabs/hardhat-solhint": "^4.0.0",

test/MaticX.ts

@@ -0,0 +1,141 @@
+import {
+	loadFixture,
+	setBalance,
+} from "@nomicfoundation/hardhat-network-helpers";
+import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
+import { expect } from "chai";
+import { ethers, upgrades } from "hardhat";
+import {
+	IERC20,
+	IFxStateRootTunnel,
+	IPolygonMigration,
+	IStakeManager,
+	MaticX,
+} from "../typechain";
+
+describe.only("MaticX", function () {
+	const stakeAmount = ethers.utils.parseUnits("100", 18);
+
+	async function deployFixture() {
+		// EOA setups
+		const maticXManager = await impersonateAccount(
+			"0x75db63125A4f04E59A1A2Ab4aCC4FC1Cd5Daddd5"
+		);
+		const maticHolder = await impersonateAccount(
+			"0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0"
+		);
+		const fxStateRootTunnelManager = await impersonateAccount(
+			"0x80A43dd35382C4919991C5Bca7f46Dd24Fde4C67"
+		);
+		const [staker] = await ethers.getSigners();
+
+		// Contract setups
+		const validatorRegistry = await ethers.getContractAt(
+			"ValidatorRegistry",
+			"0xf556442D5B77A4B0252630E15d8BbE2160870d77"
+		);
+		const stakeManager = (await ethers.getContractAt(
+			"IStakeManager",
+			"0x5e3Ef299fDDf15eAa0432E6e66473ace8c13D908"
+		)) as IStakeManager;
+
+		const matic = (await ethers.getContractAt(
+			"IERC20",
+			"0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0"
+		)) as IERC20;
+
+		const pol = (await ethers.getContractAt(
+			"IERC20",
+			"0x455e53CBB86018Ac2B8092FdCd39d8444aFFC3F6"
+		)) as IERC20;
+
+		const polygonMigration = (await ethers.getContractAt(
+			"IPolygonMigration",
+			"0x29e7DF7b6A1B2b07b731457f499E1696c60E2C4e"
+		)) as IPolygonMigration;
+
+		const fxStateRootTunnel = (await ethers.getContractAt(
+			"IFxStateRootTunnel",
+			"0x40FB804Cc07302b89EC16a9f8d040506f64dFe29",
+			fxStateRootTunnelManager
+		)) as IFxStateRootTunnel;
+
+		const MaticX = await ethers.getContractFactory("MaticX");
+		const maticX = (await upgrades.deployProxy(MaticX, [
+			validatorRegistry.address,
+			stakeManager.address,
+			matic.address,
+			maticXManager.address,
+			maticXManager.address,
+		])) as MaticX;
+
+		// Contract initializations
+		await fxStateRootTunnel
+			.connect(fxStateRootTunnelManager)
+			.setMaticX(maticX.address);
+
+		await maticX
+			.connect(maticXManager)
+			.setFxStateRootTunnel(fxStateRootTunnel.address);
+		await maticX.connect(maticXManager).setPOLToken(pol.address);
+		await maticX.connect(maticXManager).initializePOL();
+
+		// ERC20 transfers
+		await matic.connect(maticHolder).transfer(staker.address, stakeAmount);
+
+		await matic
+			.connect(maticHolder)
+			.approve(polygonMigration.address, stakeAmount);
+		await polygonMigration.connect(maticHolder).migrate(stakeAmount);
+		await pol.connect(maticHolder).transfer(staker.address, stakeAmount);
+
+		return {
+			maticX,
+			stakeManager,
+			validatorRegistry,
+			matic,
+			pol,
+			polygonMigration,
+			fxStateRootTunnel,
+			maticXManager,
+			staker,
+		};
+	}
+
+	async function impersonateAccount(
+		address: string
+	): Promise<SignerWithAddress> {
+		setBalance(address, ethers.utils.parseEther("10000"));
+		return await ethers.getImpersonatedSigner(address);
+	}
+
+	describe("Deploy the contract", function () {
+		describe("Positive", function () {
+			it("Should emit the Submit event for the Matic token", async function () {
+				const { maticX, matic, staker } =
+					await loadFixture(deployFixture);
+
+				await matic
+					.connect(staker)
+					.approve(maticX.address, stakeAmount);
+
+				const promise = maticX.connect(staker).submit(stakeAmount);
+				await expect(promise)
+					.to.emit(maticX, "Submit")
+					.withArgs(staker.address, stakeAmount);
+			});
+
+			it("Should emit the Submit event for the POL token", async function () {
+				const { maticX, pol, staker } =
+					await loadFixture(deployFixture);
+
+				await pol.connect(staker).approve(maticX.address, stakeAmount);
+
+				const promise = maticX.connect(staker).submitPOL(stakeAmount);
+				await expect(promise)
+					.to.emit(maticX, "Submit")
+					.withArgs(staker.address, stakeAmount);
+			});
+		});
+	});
+});