To ensure smooth upgradability, we should introduce a proxy contract that lets us upgrade from one implementation to another

pragma solidity ^0.8.0;

contract StakingProxy {
    uint256 public totalStaked;
    mapping(address => uint256) public stakedBalances;

    address public implementation;

    constructor(address _implementation) {
        implementation = _implementation;
    }

    fallback() external payable {
        // Forward the call to the implementation contract
        (bool success, ) = implementation.delegatecall(msg.data);
        require(success, "Delegatecall failed");
    }
}

Test

// SPDX-License-Identifier: Unlicense
pragma solidity ^0.8.13;

import "forge-std/Test.sol";

import "src/StakingContract.sol";
import "src/StakingContract2.sol";
import "src/StakingProxy.sol";

contract StakingProxyTestContract is Test {
    StakingContract wrongStakeContract;
    StakingContract2 rightStakeContract;
    StakingProxy c;

    function setUp() public {
        wrongStakeContract = new StakingContract();
        rightStakeContract = new StakingContract2();
        c = new StakingProxy(address(wrongStakeContract));
    }

    function testStake() public {
        uint value = 10 ether;
        vm.deal(0x587EFaEe4f308aB2795ca35A27Dff8c1dfAF9e3f, value);
        vm.prank(0x587EFaEe4f308aB2795ca35A27Dff8c1dfAF9e3f);
        (bool success, ) = address(c).call{value: value}(
            abi.encodeWithSignature("stake(uint256)", value)
        );
        assert(success);
        (bool success2, bytes memory data) = address(c).delegatecall(abi.encodeWithSignature("getTotalStaked()"));
        assert(success2);
        console.logBytes(data);
        uint currentStake = abi.decode(data, (uint256));
        console.log(currentStake);
        assert(currentStake == value);
    }
}

Assignment

Write the end to end test for it