From 8c3e4f6bd498d3adb6692e04902b7a28d023acc1 Mon Sep 17 00:00:00 2001
From: skosito <skostic9242@gmail.com>
Date: Wed, 7 Aug 2024 16:00:13 +0200
Subject: [PATCH] add more connector tests

---
 v2/test/ZetaConnectorNative.t.sol    | 37 +++++++++++++++++++++++++++
 v2/test/ZetaConnectorNonNative.t.sol | 38 ++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/v2/test/ZetaConnectorNative.t.sol b/v2/test/ZetaConnectorNative.t.sol
index 87c0c1d6..aaabc455 100644
--- a/v2/test/ZetaConnectorNative.t.sol
+++ b/v2/test/ZetaConnectorNative.t.sol
@@ -40,9 +40,11 @@ contract ZetaConnectorNativeTest is
     address destination;
     address tssAddress;
 
+    error EnforcedPause();
     error AccessControlUnauthorizedAccount(address account, bytes32 neededRole);
 
     bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
+    bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE");
 
     function setUp() public {
         owner = address(this);
@@ -88,6 +90,41 @@ contract ZetaConnectorNativeTest is
         assertEq(balanceAfter, amount);
     }
 
+    function testWithdrawTogglePause() public {
+        uint256 amount = 100_000;
+        bytes32 internalSendHash = "";
+
+        vm.prank(tssAddress);
+        vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, tssAddress, PAUSER_ROLE));
+        zetaConnector.pause();
+
+        vm.prank(tssAddress);
+        vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, tssAddress, PAUSER_ROLE));
+        zetaConnector.unpause();
+
+        vm.prank(owner);
+        zetaConnector.pause();
+
+        vm.expectRevert(EnforcedPause.selector);
+        vm.prank(tssAddress);
+        zetaConnector.withdraw(destination, amount, internalSendHash);
+
+        vm.prank(owner);
+        zetaConnector.unpause();
+
+        uint256 balanceBefore = zetaToken.balanceOf(destination);
+        assertEq(balanceBefore, 0);
+
+        bytes memory data = abi.encodeWithSignature("transfer(address,uint256)", destination, amount);
+        vm.expectCall(address(zetaToken), 0, data);
+        vm.expectEmit(true, true, true, true, address(zetaConnector));
+        emit Withdraw(destination, amount);
+        vm.prank(tssAddress);
+        zetaConnector.withdraw(destination, amount, internalSendHash);
+        uint256 balanceAfter = zetaToken.balanceOf(destination);
+        assertEq(balanceAfter, amount);
+    }
+
     function testWithdrawFailsIfSenderIsNotWithdrawer() public {
         uint256 amount = 100_000;
         bytes32 internalSendHash = "";
diff --git a/v2/test/ZetaConnectorNonNative.t.sol b/v2/test/ZetaConnectorNonNative.t.sol
index 9397459d..11dc3a31 100644
--- a/v2/test/ZetaConnectorNonNative.t.sol
+++ b/v2/test/ZetaConnectorNonNative.t.sol
@@ -42,9 +42,11 @@ contract ZetaConnectorNonNativeTest is
 
     error AccessControlUnauthorizedAccount(address account, bytes32 neededRole);
     error ExceedsMaxSupply();
+    error EnforcedPause();
 
     bytes32 public constant WITHDRAWER_ROLE = keccak256("WITHDRAWER_ROLE");
     bytes32 public constant TSS_ROLE = keccak256("TSS_ROLE");
+    bytes32 public constant PAUSER_ROLE = keccak256("PAUSER_ROLE");
 
     function setUp() public {
         owner = address(this);
@@ -91,6 +93,42 @@ contract ZetaConnectorNonNativeTest is
         assertEq(balanceAfter, amount);
     }
 
+    function testWithdrawTogglePause() public {
+        uint256 amount = 100_000;
+        bytes32 internalSendHash = "";
+
+        vm.prank(tssAddress);
+        vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, tssAddress, PAUSER_ROLE));
+        zetaConnector.pause();
+
+        vm.prank(tssAddress);
+        vm.expectRevert(abi.encodeWithSelector(AccessControlUnauthorizedAccount.selector, tssAddress, PAUSER_ROLE));
+        zetaConnector.unpause();
+
+        vm.prank(owner);
+        zetaConnector.pause();
+
+        vm.expectRevert(EnforcedPause.selector);
+        vm.prank(tssAddress);
+        zetaConnector.withdraw(destination, amount, internalSendHash);
+
+        vm.prank(owner);
+        zetaConnector.unpause();
+
+        uint256 balanceBefore = zetaToken.balanceOf(destination);
+        assertEq(balanceBefore, 0);
+
+        bytes memory data =
+            abi.encodeWithSignature("mint(address,uint256,bytes32)", destination, amount, internalSendHash);
+        vm.expectCall(address(zetaToken), 0, data);
+        vm.expectEmit(true, true, true, true, address(zetaConnector));
+        emit Withdraw(destination, amount);
+        vm.prank(tssAddress);
+        zetaConnector.withdraw(destination, amount, internalSendHash);
+        uint256 balanceAfter = zetaToken.balanceOf(destination);
+        assertEq(balanceAfter, amount);
+    }
+
     function testWithdrawFailsIfSenderIsNotWithdrawer() public {
         uint256 amount = 100_000;
         bytes32 internalSendHash = "";