update existing tests

nori-dot-eco · Mar 8, 2024 · 6bcc5d9 · 6bcc5d9
1 parent 599f58e
commit 6bcc5d9
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 31 deletions.
diff --git a/contracts/Market.sol b/contracts/Market.sol
@@ -781,19 +781,22 @@ contract Market is
   }
 
   /**
-   * @notice Withdraws a removal to the supplier.
-   * @dev Withdraws a removal to the supplier address encoded in the removal ID.
+   * @notice Withdraws a removal from the Market.
+   * @dev Withdraws a removal to the specified address, which much be the supplier of the removal
+   * or have the `Removal.CONSIGNOR_ROLE`.
    *
    * ##### Requirements:
    *
    * - Can only be used when this contract is not paused.
+   * - Can only withdraw to the supplier of the removal or an address with the `Removal.CONSIGNOR_ROLE`.
    * @param removalId The ID of the removal to withdraw from the market.
+   * @param to The address to which the removal will be transferred.
    */
-  function withdraw(uint256 removalId) external whenNotPaused {
+  function withdraw(uint256 removalId, address to) external whenNotPaused {
     address supplierAddress = RemovalIdLib.supplierAddress({
       removalId: removalId
     });
-    if (_isAuthorizedWithdrawal({owner: supplierAddress})) {
+    if (_isAuthorizedWithdrawal({supplier: supplierAddress, to: to})) {
       _removeActiveRemoval({
         removalId: removalId,
         supplierAddress: supplierAddress
@@ -1349,16 +1352,26 @@ contract Market is
   }
 
   /**
-   * @dev Authorizes withdrawal for the removal. Reverts if the caller is not the owner of the removal and
-   * does not have the role `MARKET_ADMIN_ROLE`.
-   * @param owner The owner of the removal.
-   * @return Returns true if the caller is the owner, an approved spender, or has the role `MARKET_ADMIN_ROLE`,
-   * false otherwise.
+   * @dev Authorizes withdrawal for the removal. Reverts if the caller is not the supplier of the removal and
+   * does not have the role `MARKET_ADMIN_ROLE`, or if the recipient of the removal is either not the supplier
+   * of the removal or does not have the `Removal.CONSIGNOR_ROLE`.
+   * @param supplier The supplier of the removal.
+   * @return Returns true if the caller is the supplier, an approved spender, or has the role `MARKET_ADMIN_ROLE`,
+   * and the recipient is the removal supplier or has the `Removal.CONSIGNOR_ROLE`, false otherwise.
    */
-  function _isAuthorizedWithdrawal(address owner) internal view returns (bool) {
-    return (_msgSender() == owner ||
-      hasRole({role: MARKET_ADMIN_ROLE, account: _msgSender()}) ||
-      _removal.isApprovedForAll({account: owner, operator: _msgSender()}));
+  function _isAuthorizedWithdrawal(
+    address supplier,
+    address to
+  ) internal view returns (bool) {
+    return
+      (_msgSender() == supplier ||
+        hasRole({role: MARKET_ADMIN_ROLE, account: _msgSender()}) ||
+        _removal.isApprovedForAll({
+          account: supplier,
+          operator: _msgSender()
+        })) &&
+      (_removal.hasRole({role: _removal.CONSIGNOR_ROLE(), account: to}) ||
+        supplier == to);
   }
 
   /**

diff --git a/test/Market.t.sol b/test/Market.t.sol
@@ -408,7 +408,7 @@ contract Market_withdraw_as_supplier is MarketBalanceTestHelper {
 
   function test() external {
     vm.prank(_namedAccounts.supplier);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [_amountPerRemoval];
     _expectedMarketSupply = 0;
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -441,7 +441,7 @@ contract Market_withdraw_as_operator is MarketBalanceTestHelper {
 
   function test() external {
     vm.prank(_namedAccounts.supplier2);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [_amountPerRemoval];
     _expectedMarketSupply = 0;
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -466,7 +466,7 @@ contract Market_withdraw_as_DEFAULT_ADMIN_ROLE is MarketBalanceTestHelper {
   }
 
   function test() external {
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [_amountPerRemoval];
     _expectedMarketSupply = 0;
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -493,7 +493,7 @@ contract Market_withdraw_reverts is MarketBalanceTestHelper {
   function test() external {
     vm.prank(_namedAccounts.supplier2);
     vm.expectRevert(UnauthorizedWithdrawal.selector);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier2});
     _assertCorrectStates();
   }
 }
@@ -515,7 +515,7 @@ contract Market_withdraw_1x3_center is MarketBalanceTestHelper {
 
   function test() external {
     vm.prank(_namedAccounts.supplier);
-    _market.withdraw(_removalIds[1]);
+    _market.withdraw({removalId: _removalIds[1], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [0, _amountPerRemoval, 0];
     _expectedMarketSupply = _amountPerRemoval * (_removalIds.length - 1);
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -541,7 +541,7 @@ contract Market_withdraw_2x1_front is MarketBalanceTestHelper {
 
   function test() external {
     vm.prank(_namedAccounts.supplier);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [_amountPerRemoval, 0];
     _expectedMarketSupply = _amountPerRemoval * (_removalIds.length - 1);
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -568,7 +568,7 @@ contract Market_withdraw_2x1_front_relist is MarketBalanceTestHelper {
     ];
     _assertListedState();
     vm.prank(_namedAccounts.supplier);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     _expectedRemovalBalances = [_amountPerRemoval, 0];
     _expectedMarketSupply = _amountPerRemoval * (_removalIds.length - 1);
     _expectedTokenCount.set(_namedAccounts.supplier, 1);
@@ -604,7 +604,7 @@ contract Market_withdraw_2x1_back is MarketBalanceTestHelper {
 
   function test() external {
     vm.prank(_namedAccounts.supplier2);
-    _market.withdraw(_removalIds[1]);
+    _market.withdraw({removalId: _removalIds[1], to: _namedAccounts.supplier2});
     _expectedRemovalBalances = [0, _amountPerRemoval];
     _expectedMarketSupply = _amountPerRemoval * (_removalIds.length - 1);
     _expectedTokenCount.set(_namedAccounts.supplier, 0);
@@ -629,31 +629,50 @@ contract Market_SANCTION_ALLOWLIST_ROLE is UpgradeableMarket {
   }
 }
 
-contract Market__isAuthorizedWithdrawal_true is NonUpgradeableMarket {
+contract Market__isAuthorizedWithdrawal_true is
+  NonUpgradeableMarket,
+  UpgradeableRemoval
+{
+  Removal private _mockRemoval;
+
   function setUp() external {
+    _mockRemoval = _deployRemoval();
+
     vm.store(
       address(this),
-      bytes32(uint256(301)), // sets the _removal storage slot to the market contract to enable mock calls
-      bytes32(uint256(uint160(address(this))))
+      bytes32(uint256(301)), // sets the _removal storage slot in the market contract to enable mock calls
+      bytes32(uint256(uint160(address(_mockRemoval))))
     );
   }
 
   function test_returnsTrueWhenMsgSenderEqualsOwner() external {
-    assertEq(_isAuthorizedWithdrawal({owner: _msgSender()}), true);
+    assertEq(
+      _isAuthorizedWithdrawal({supplier: _msgSender(), to: _msgSender()}),
+      true
+    );
   }
 
   function test_returnsTrueWhenMsgSenderHasDefaultAdminRole() external {
     _grantRole({role: MARKET_ADMIN_ROLE, account: _msgSender()});
-    assertEq(_isAuthorizedWithdrawal({owner: _namedAccounts.supplier}), true);
+    assertEq(
+      _isAuthorizedWithdrawal({
+        supplier: _namedAccounts.supplier,
+        to: _namedAccounts.supplier
+      }),
+      true
+    );
   }
 
   function test_returnsTrueWhenMsgSenderIsApprovedForAll() external {
     vm.mockCall(
-      address(this),
-      abi.encodeWithSelector(IERC1155Upgradeable.isApprovedForAll.selector),
+      address(_mockRemoval),
+      abi.encodeWithSelector(_mockRemoval.isApprovedForAll.selector),
       abi.encode(true)
     );
-    assertEq(_isAuthorizedWithdrawal({owner: address(0)}), true);
+    assertEq(
+      _isAuthorizedWithdrawal({supplier: address(1), to: address(1)}),
+      true
+    );
   }
 }
 
@@ -672,7 +691,13 @@ contract Market__isAuthorizedWithdrawal_false is NonUpgradeableMarket {
   }
 
   function test_returnsFalseWhenAllConditionsAreFalse() external {
-    assertEq(_isAuthorizedWithdrawal({owner: _namedAccounts.supplier}), false);
+    assertEq(
+      _isAuthorizedWithdrawal({
+        supplier: _namedAccounts.supplier,
+        to: _namedAccounts.supplier
+      }),
+      false
+    );
   }
 }
 

diff --git a/test/Removal.t.sol b/test/Removal.t.sol
@@ -1149,7 +1149,7 @@ contract Removal_getMarketBalance is UpgradeableMarket {
       signedPermit.s
     );
     assertEq(_removal.getMarketBalance(), amountToList - amountToSell);
-    _market.withdraw(_removalIds[0]);
+    _market.withdraw({removalId: _removalIds[0], to: _namedAccounts.supplier});
     assertEq(_removal.getMarketBalance(), 0);
   }
 }