From 39bd9bd63ffb285e67ba3d37497c11370a0a5f38 Mon Sep 17 00:00:00 2001
From: Zodomo <zodomo@proton.me>
Date: Tue, 26 Nov 2024 17:34:49 -0600
Subject: [PATCH] add AVS upgrade tests

---
 contracts/avs/.gas-snapshot              |  4 +-
 contracts/avs/test/OmniAVS_upgrade.t.sol | 63 ++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 contracts/avs/test/OmniAVS_upgrade.t.sol

diff --git a/contracts/avs/.gas-snapshot b/contracts/avs/.gas-snapshot
index 052caf21a..d65a06f1e 100644
--- a/contracts/avs/.gas-snapshot
+++ b/contracts/avs/.gas-snapshot
@@ -59,4 +59,6 @@ OmniAVS_canRegister_Test:test_canRegister_notAllowed() (gas: 225902)
 OmniAVS_canRegister_Test:test_canRegister_notOperator() (gas: 120672)
 OmniAVS_initialize_Test:test_initialize_defaultParams_succeeds() (gas: 3606316)
 OmniAVS_syncWithOmni_Test:test_depositBeaconEth_succeeds() (gas: 638274)
-OmniAVS_syncWithOmni_Test:test_unsupportedStrategyDeposit_succeeds() (gas: 1677678)
\ No newline at end of file
+OmniAVS_syncWithOmni_Test:test_unsupportedStrategyDeposit_succeeds() (gas: 1677678)
+OmniAVS_upgrade_Test:test_upgrade_invalidProxyOwner_reverts() (gas: 2698429)
+OmniAVS_upgrade_Test:test_upgrade_succeeds() (gas: 2756363)
\ No newline at end of file
diff --git a/contracts/avs/test/OmniAVS_upgrade.t.sol b/contracts/avs/test/OmniAVS_upgrade.t.sol
new file mode 100644
index 000000000..25c9864ef
--- /dev/null
+++ b/contracts/avs/test/OmniAVS_upgrade.t.sol
@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: GPL-3.0-only
+pragma solidity =0.8.12;
+
+import { Base } from "./common/Base.sol";
+
+import { OmniAVS } from "src/OmniAVS.sol";
+
+import { IDelegationManager } from "src/ext/IDelegationManager.sol";
+import { IAVSDirectory } from "eigenlayer-contracts/src/contracts/interfaces/IAVSDirectory.sol";
+import { ITransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+
+contract OmniAVS_upgrade_Test is Base {
+    function _checkAndHashState() internal view returns (bytes32) {
+        address _avsDirectory = omniAVS.avsDirectory();
+        uint32 _maxOperatorCount = omniAVS.maxOperatorCount();
+        uint64 _omniChainId = omniAVS.omniChainId();
+        uint64 _xcallGasLimitPerOperator = omniAVS.xcallGasLimitPerOperator();
+        uint64 _xcallBaseGasLimit = omniAVS.xcallBaseGasLimit();
+        uint96 _minOperatorStake = omniAVS.minOperatorStake();
+        bool _allowlistEnabled = omniAVS.allowlistEnabled();
+        address _ethStakeInbox = omniAVS.ethStakeInbox();
+        address _omni = address(omniAVS.omni());
+
+        return keccak256(
+            abi.encode(
+                _avsDirectory,
+                _maxOperatorCount,
+                _omniChainId,
+                _xcallGasLimitPerOperator,
+                _xcallBaseGasLimit,
+                _minOperatorStake,
+                _allowlistEnabled,
+                _ethStakeInbox,
+                _omni
+            )
+        );
+    }
+
+    function test_upgrade_succeeds() public {
+        address impl =
+            address(new OmniAVS(IDelegationManager(address(delegation)), IAVSDirectory(address(avsDirectory))));
+        ITransparentUpgradeableProxy proxy = ITransparentUpgradeableProxy(address(omniAVS));
+
+        bytes32 beforeUpgradeHash = _checkAndHashState();
+
+        vm.prank(proxyAdminOwner);
+        proxyAdmin.upgrade(proxy, impl);
+
+        bytes32 afterUpgradeHash = _checkAndHashState();
+
+        assertEq(proxyAdmin.getProxyImplementation(proxy), impl);
+        assertEq(beforeUpgradeHash, afterUpgradeHash);
+    }
+
+    function test_upgrade_invalidProxyOwner_reverts() public {
+        address impl =
+            address(new OmniAVS(IDelegationManager(address(delegation)), IAVSDirectory(address(avsDirectory))));
+        ITransparentUpgradeableProxy proxy = ITransparentUpgradeableProxy(address(omniAVS));
+
+        vm.expectRevert("Ownable: caller is not the owner");
+        proxyAdmin.upgrade(proxy, impl);
+    }
+}