feat: adding call tokens expectation

defi-wonderland · Jul 18, 2024 · ffb1bc0 · ffb1bc0
1 parent b2056be
commit ffb1bc0
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 1 deletion.
diff --git a/src/contracts/BCoWHelper.sol b/src/contracts/BCoWHelper.sol
@@ -104,7 +104,7 @@ contract BCoWHelper is ICOWAMMPoolHelper, BMath {
   }
 
   /// @inheritdoc ICOWAMMPoolHelper
-  function tokens(address pool) public view returns (address[] memory tokens_) {
+  function tokens(address pool) public view virtual returns (address[] memory tokens_) {
     // reverts in case pool is not deployed by the helper's factory
     if (!IBCoWFactory(factory).isBPool(pool)) {
       revert PoolDoesNotExist();

diff --git a/test/manual-smock/MockBCoWHelper.sol b/test/manual-smock/MockBCoWHelper.sol
@@ -20,6 +20,19 @@ contract MockBCoWHelper is BCoWHelper, Test {
     return _APP_DATA;
   }
 
+  // NOTE: manually added method (public overrides not supported in smock)
+  function tokens(address pool) public view override returns (address[] memory tokens_) {
+    (bool _success, bytes memory _data) = address(this).staticcall(abi.encodeWithSignature('tokens(address)', pool));
+
+    if (_success) return abi.decode(_data, (address[]));
+    else return super.tokens(pool);
+  }
+
+  // NOTE: manually added method (public overrides not supported in smock)
+  function expectCall_tokens(address pool) public {
+    vm.expectCall(address(this), abi.encodeWithSignature('tokens(address)', pool));
+  }
+
   // BCoWHelper methods
   constructor(address factory_) BCoWHelper(factory_) {}
 

diff --git a/test/unit/BCoWHelper.t.sol b/test/unit/BCoWHelper.t.sol
@@ -120,6 +120,10 @@ contract BCoWHelperTest is Test {
   }
 
   function test_OrderWhenThePoolIsSupported(bytes32 domainSeparator) external {
+    // it should call tokens
+    helper.mock_call_tokens(address(pool), tokens);
+    helper.expectCall_tokens(address(pool));
+
     // it should query the domain separator from the pool
     pool.expectCall_SOLUTION_SETTLER_DOMAIN_SEPARATOR();
     pool.mock_call_SOLUTION_SETTLER_DOMAIN_SEPARATOR(domainSeparator);

diff --git a/test/unit/BCoWHelper.tree b/test/unit/BCoWHelper.tree
@@ -19,6 +19,7 @@ BCoWHelperTest::order
 ├── when the pool is not supported
 │   └── it should revert
 ├── when the pool is supported
+│   ├── it should call tokens
 │   ├── it should query the domain separator from the pool
 │   ├── it should return a valid pool order
 │   ├── it should return a commit pre-interaction