I clean up and add additional tests for transfers.

defientco · Aug 6, 2023 · 308b18c · 308b18c
1 parent b742345
commit 308b18c
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 25 deletions.
diff --git a/test/hooks/AfterTokenTransfers.t.sol b/test/hooks/AfterTokenTransfers.t.sol
@@ -33,13 +33,8 @@ contract AfterTokenTransfersHookTest is DSTest, HookUtils {
         vm.expectRevert();
         erc721Mock.setHook(AfterTokenTransfers, address(hookMock));
 
-        vm.prank(DEFAULT_OWNER_ADDRESS);
-        erc721Mock.setHook(AfterTokenTransfers, address(hookMock));
-
-        assertEq(
-            address(hookMock),
-            address(erc721Mock.getHook(AfterTokenTransfers))
-        );
+        //  Verify Success
+        _setHook(address(erc721Mock), AfterTokenTransfers, address(hookMock));
     }
 
     function test_afterTokenTransfersHook(
@@ -76,4 +71,26 @@ contract AfterTokenTransfersHookTest is DSTest, HookUtils {
                 .selector
         );
     }
+
+    function test_turn_off_hook(
+        address _firstOwner,
+        address _secondOwner,
+        uint256 startTokenId,
+        uint256 quantity
+    ) public {
+        test_afterTokenTransfersHook(
+            _firstOwner,
+            _secondOwner,
+            startTokenId,
+            quantity
+        );
+
+        _setHook(address(erc721Mock), AfterTokenTransfers, address(0));
+        _assertNormalTransfer(
+            address(erc721Mock),
+            _secondOwner,
+            _firstOwner,
+            startTokenId
+        );
+    }
 }
diff --git a/test/hooks/BeforeTokenTransfers.t.sol b/test/hooks/BeforeTokenTransfers.t.sol
@@ -30,12 +30,8 @@ contract BeforeTokenTransfersHookTest is DSTest, HookUtils {
         vm.expectRevert();
         erc721Mock.setHook(BeforeTokenTransfers, address(hookMock));
 
-        vm.prank(DEFAULT_OWNER_ADDRESS);
-        erc721Mock.setHook(BeforeTokenTransfers, address(hookMock));
-        assertEq(
-            address(hookMock),
-            address(erc721Mock.getHook(BeforeTokenTransfers))
-        );
+        // set hook
+        _setHook(address(erc721Mock), BeforeTokenTransfers, address(hookMock));
     }
 
     function test_beforeTokenTransfersHook(
@@ -73,17 +69,25 @@ contract BeforeTokenTransfersHookTest is DSTest, HookUtils {
         );
     }
 
-    function _assertHookRevert(
-        address _from,
-        address _to,
-        uint256 _tokenId
-    ) internal {
-        vm.prank(_from);
-        vm.expectRevert(
-            BeforeTokenTransfersHookMock
-                .BeforeTokenTransfersHook_Executed
-                .selector
+    function test_turn_off_hook(
+        address _firstOwner,
+        address _secondOwner,
+        uint256 startTokenId,
+        uint256 quantity
+    ) public {
+        test_beforeTokenTransfersHook(
+            _firstOwner,
+            _secondOwner,
+            startTokenId,
+            quantity
+        );
+
+        _setHook(address(erc721Mock), BeforeTokenTransfers, address(0));
+        _assertNormalTransfer(
+            address(erc721Mock),
+            _secondOwner,
+            _firstOwner,
+            startTokenId
         );
-        erc721Mock.transferFrom(_from, _to, _tokenId);
     }
 }
diff --git a/test/utils/HookUtils.sol b/test/utils/HookUtils.sol
@@ -5,7 +5,6 @@ import {Vm} from "forge-std/Vm.sol";
 import {DSTest} from "ds-test/test.sol";
 import {ERC721ACHMock} from "./ERC721ACHMock.sol";
 import {IERC721A} from "lib/ERC721A/contracts/IERC721A.sol";
-import {BeforeTokenTransfersHookMock} from "../utils/hooks/BeforeTokenTransfersHookMock.sol";
 import {IERC721ACH} from "../../src/interfaces/IERC721ACH.sol";
 
 contract HookUtils is DSTest {
@@ -53,4 +52,17 @@ contract HookUtils is DSTest {
         vm.expectRevert(_err);
         ERC721ACHMock(_target).transferFrom(_from, _to, _tokenId);
     }
+
+    function _setHook(
+        address _target,
+        IERC721ACH.HookType _type,
+        address _hook
+    ) internal {
+        vm.prank(DEFAULT_OWNER_ADDRESS);
+        ERC721ACHMock(_target).setHook(_type, address(_hook));
+        assertEq(
+            address(_hook),
+            address(ERC721ACHMock(_target).getHook(_type))
+        );
+    }
 }