Add testing for HybridProxy

Author: Amxx

contracts/mocks/UpgradeableImplementation.sol

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.27;
+
+import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
+
+// Note using UUPSUpgradeable here will be too restrictive
+// - when upgrading an HydridProxy from beacon mode to UUPS mode, the `_checkProxy` check fails (implementation slot doesn't contain "self")
+// - when upgrading an HydridProxy from UUPS mode to beacon mode, the IERC1822Proxiable check fails (beacon doesn't implement it)
+// So we manually implement `upgradeToAndCall` without any checks
+contract UpgradeableImplementationMock {
+    address private immutable __self = address(this);
+    uint256 public immutable version;
+
+    error UnexpectedCall(address, uint256, bytes);
+
+    constructor(uint256 _version) {
+        version = _version;
+    }
+
+    function upgradeToAndCall(address newImplementation, bytes memory data) public payable virtual {
+        ERC1967Utils.upgradeToAndCall(newImplementation, data);
+    }
+
+    receive() external payable {
+        revert UnexpectedCall(msg.sender, msg.value, "0x");
+    }
+
+    fallback() external payable {
+        revert UnexpectedCall(msg.sender, msg.value, msg.data);
+    }
+}

contracts/mocks/import.sol

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.20;
+
+import {UpgradeableBeacon} from "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol";

contracts/proxy/HybridProxy.sol

@@ -11,7 +11,9 @@ import {Address} from "@openzeppelin/contracts/utils/Address.sol";
  * @dev A version of an ERC-1967 proxy that uses the address stored in the implementation slot as a beacon.
  *
  * The design allows to set an initial beacon that the contract may quit by upgrading to its own implementation
- * afterwards.
+ * afterwards. Transition between the "beacon mode" and the "direct mode" require implementation that expose an
+ * upgrade mechanism that writes to the ERC-1967 implementation slot. Note that UUPSUpgradable includes security
+ * checks that are not compatible with this proxy design.
  *
  * WARNING: The fallback mechanism relies on the implementation not to define the {IBeacon-implementation} function.
  * Consider that if your implementation has this function, it'll be assumed as the beacon address, meaning that

test/proxy/HybridProxy.test.js

@@ -0,0 +1,52 @@
+const { ethers } = require('hardhat');
+const { expect } = require('chai');
+const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
+
+async function fixture() {
+  const [admin] = await ethers.getSigners();
+
+  const implementation1 = await ethers.deployContract('UpgradeableImplementationMock', [1]);
+  const implementation2 = await ethers.deployContract('UpgradeableImplementationMock', [2]);
+  const implementation3 = await ethers.deployContract('UpgradeableImplementationMock', [3]);
+  const beacon = await ethers.deployContract('UpgradeableBeacon', [implementation1, admin]);
+  const proxy = await ethers
+    .deployContract('HybridProxy', [beacon, '0x'])
+    .then(({ target }) => implementation1.attach(target));
+
+  return { admin, beacon, proxy, implementation1, implementation2, implementation3 };
+}
+
+describe('HybridProxy', function () {
+  beforeEach(async function () {
+    Object.assign(this, await loadFixture(fixture));
+  });
+
+  it('setup at construction', async function () {
+    const data = ethers.randomBytes(128);
+    await expect(ethers.deployContract('HybridProxy', [this.beacon, data]))
+      .to.be.revertedWithCustomError(this.implementation1, 'UnexpectedCall')
+      .withArgs(this.admin, 0, data);
+  });
+
+  it('forwards calls', async function () {
+    await expect(this.implementation1.attach(this.proxy).version()).to.eventually.equal(1);
+  });
+
+  it('beacon upgrade', async function () {
+    await this.beacon.upgradeTo(this.implementation2);
+    await expect(this.implementation1.attach(this.proxy).version()).to.eventually.equal(2);
+  });
+
+  it('decouple/recouple', async function () {
+    await this.proxy.upgradeToAndCall(this.implementation3, '0x');
+    await expect(this.implementation1.attach(this.proxy).version()).to.eventually.equal(3);
+
+    // beacon updated no longer affect the upgrade process
+    await this.beacon.upgradeTo(this.implementation2);
+    await expect(this.implementation1.attach(this.proxy).version()).to.eventually.equal(3);
+
+    // recouple to beacon
+    await this.proxy.upgradeToAndCall(this.beacon, '0x');
+    await expect(this.implementation1.attach(this.proxy).version()).to.eventually.equal(2);
+  });
+});