Merge pull request #859 from gitcoindev/test-diamond-cut

test: facetCut tests

packages/contracts/src/dollar/mocks/MockFacet.sol

@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+import {LibDiamond} from "../libraries/LibDiamond.sol";
+
+interface IMockFacet {
+    function functionA() external returns (uint256);
+    function functionB() external returns (uint256);
+}
+
+struct NewSimpleStorage {
+    uint256 slot1;
+}
+
+struct NewStorageExtended {
+    uint256 slot1;
+    uint256 slot2;
+}
+
+bytes32 constant NEW_STORAGE_POSITION = keccak256(
+    "ubiquity.contracts.mock.storage"
+);
+
+contract MockFacetInitializer {
+    function initialize() external {
+        NewStorageExtended storage newStorage;
+        bytes32 position = NEW_STORAGE_POSITION;
+        assembly {
+            newStorage.slot := position
+        }
+        newStorage.slot1 = 2;
+        newStorage.slot2 = 22;
+    }
+
+    function initializeRevert() external pure returns (uint256) {
+        revert();
+    }
+
+    function initializeRevertWithMessage() external pure returns (uint256) {
+        revert("MockFacetError");
+    }
+}
+
+contract MockFacetWithPureFunctions is IMockFacet {
+    function functionA() external pure returns (uint256) {
+        return 1;
+    }
+
+    function functionB() external pure returns (uint256) {
+        return 2;
+    }
+}
+
+contract MockFacetWithStorageWriteFunctions is IMockFacet {
+    function functionA() external view returns (uint256) {
+        NewSimpleStorage storage newStorage;
+        bytes32 position = NEW_STORAGE_POSITION;
+        assembly {
+            newStorage.slot := position
+        }
+        return newStorage.slot1;
+    }
+
+    function functionB() external returns (uint256) {
+        NewSimpleStorage storage newStorage;
+        bytes32 position = NEW_STORAGE_POSITION;
+        assembly {
+            newStorage.slot := position
+        }
+        newStorage.slot1 = 1;
+        return newStorage.slot1;
+    }
+}
+
+contract MockFacetWithExtendedStorageWriteFunctions is IMockFacet {
+    function functionA() external view returns (uint256) {
+        NewStorageExtended storage newStorage;
+        bytes32 position = NEW_STORAGE_POSITION;
+        assembly {
+            newStorage.slot := position
+        }
+        return newStorage.slot1;
+    }
+
+    function functionB() external view returns (uint256) {
+        NewStorageExtended storage newStorage;
+        bytes32 position = NEW_STORAGE_POSITION;
+        assembly {
+            newStorage.slot := position
+        }
+        return newStorage.slot2;
+    }
+}

packages/contracts/test/diamond/DiamondTest.t.sol

@@ -2,8 +2,15 @@
 pragma solidity 0.8.19;
 
 import "./DiamondTestSetup.sol";
+import {IMockFacet, MockFacetInitializer, MockFacetWithPureFunctions, MockFacetWithStorageWriteFunctions, MockFacetWithExtendedStorageWriteFunctions} from "../../src/dollar/mocks/MockFacet.sol";
 
 contract TestDiamond is DiamondTestSetup {
+    address pureFacet = address(new MockFacetWithPureFunctions());
+    address writeFacet = address(new MockFacetWithStorageWriteFunctions());
+    address writeFacetWithInitializer =
+        address(new MockFacetWithExtendedStorageWriteFunctions());
+    address facetInitializer = address(new MockFacetInitializer());
+
     function test_ShouldSupportInspectingFacetsAndFunctions() public {
         bool isSupported = IERC165(address(diamond)).supportsInterface(
             type(IDiamondLoupe).interfaceId
@@ -48,6 +55,238 @@ contract TestDiamond is DiamondTestSetup {
         }
     }
 
+    function testCutFacetShouldRevertWhenNotOwner() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: selectorsOfCollectableDustFacet
+        });
+
+        vm.expectRevert("LibDiamond: Must be contract owner");
+
+        vm.prank(user1);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+    }
+
+    function testCutFacetShouldRevertWhenFunctionAlreadyExists() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(collectableDustFacetImplementation),
+            action: FacetCutAction.Add,
+            functionSelectors: selectorsOfCollectableDustFacet
+        });
+
+        vm.expectRevert(
+            "LibDiamondCut: Can't add function that already exists"
+        );
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+    }
+
+    function testCutFacetShouldRevertWhenNoSelectorsProvidedForFacetForCut()
+        public
+    {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: new bytes4[](0)
+        });
+
+        vm.expectRevert("LibDiamondCut: No selectors in facet to cut");
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+    }
+
+    function testCutFacetShouldRevertWhenAddToZeroAddress() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(0),
+            action: FacetCutAction.Add,
+            functionSelectors: selectorsOfCollectableDustFacet
+        });
+
+        vm.expectRevert("LibDiamondCut: Add facet can't be address(0)");
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+    }
+
+    function testCutFacetShouldRevertWhenFacetInitializerReverts() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Replace,
+            functionSelectors: selectorsOfCollectableDustFacet
+        });
+
+        vm.expectRevert();
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(
+            facetCut,
+            facetInitializer,
+            abi.encodeWithSelector(
+                MockFacetInitializer.initializeRevert.selector
+            )
+        );
+    }
+
+    function testCutFacetShouldRevertWithMessageWhenFacetInitializerWithMessageReverts()
+        public
+    {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Replace,
+            functionSelectors: selectorsOfCollectableDustFacet
+        });
+
+        vm.expectRevert("MockFacetError");
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(
+            facetCut,
+            facetInitializer,
+            abi.encodeWithSelector(
+                MockFacetInitializer.initializeRevertWithMessage.selector
+            )
+        );
+    }
+
+    function testCutFacetAddSimplePureFacet() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = MockFacetWithPureFunctions.functionA.selector;
+        selectors[1] = MockFacetWithPureFunctions.functionB.selector;
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 1);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 2);
+    }
+
+    function testCutFacetAddSimpleWriteFacet() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = MockFacetWithStorageWriteFunctions.functionA.selector;
+        selectors[1] = MockFacetWithStorageWriteFunctions.functionB.selector;
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(writeFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 0);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 1);
+    }
+
+    function testCutFacetAddWriteFacetWithInitializer() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = MockFacetWithExtendedStorageWriteFunctions
+            .functionA
+            .selector;
+        selectors[1] = MockFacetWithExtendedStorageWriteFunctions
+            .functionB
+            .selector;
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(writeFacetWithInitializer),
+            action: FacetCutAction.Add,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(
+            facetCut,
+            facetInitializer,
+            abi.encodeWithSelector(MockFacetInitializer.initialize.selector)
+        );
+
+        // initializer should set 2 and 22 values
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 2);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 22);
+    }
+
+    function testCutFacetReplaceFacet() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = MockFacetWithPureFunctions.functionA.selector;
+        selectors[1] = MockFacetWithPureFunctions.functionB.selector;
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(pureFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 1);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 2);
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(writeFacet),
+            action: FacetCutAction.Replace,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 0);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 1);
+    }
+
+    function testCutFacetRemoveFacetFunctions() public {
+        FacetCut[] memory facetCut = new FacetCut[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = MockFacetWithStorageWriteFunctions.functionA.selector;
+        selectors[1] = MockFacetWithStorageWriteFunctions.functionB.selector;
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(writeFacet),
+            action: FacetCutAction.Add,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        assertEq(IMockFacet(address(diamondCutFacet)).functionA(), 0);
+        assertEq(IMockFacet(address(diamondCutFacet)).functionB(), 1);
+
+        facetCut[0] = FacetCut({
+            facetAddress: address(0),
+            action: FacetCutAction.Remove,
+            functionSelectors: selectors
+        });
+
+        vm.prank(owner);
+        diamondCutFacet.diamondCut(facetCut, address(0x0), "");
+
+        vm.expectRevert("Diamond: Function does not exist");
+        IMockFacet(address(diamondCutFacet)).functionA();
+        vm.expectRevert("Diamond: Function does not exist");
+        IMockFacet(address(diamondCutFacet)).functionB();
+    }
+
     function testSelectors_ShouldBeAssociatedWithCorrectFacet() public {
         for (uint256 i; i < facetAddressList.length; i++) {
             if (compareStrings(facetNames[i], "DiamondCutFacet")) {