From ffb1bc0e4939360a7fb66bc57a5ab5a126d1d359 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wei=C3=9Fer=20Hase?= <wei3erHase@protonmail.com>
Date: Thu, 18 Jul 2024 12:10:43 +0200
Subject: [PATCH] feat: adding call tokens expectation

---
 src/contracts/BCoWHelper.sol         |  2 +-
 test/manual-smock/MockBCoWHelper.sol | 13 +++++++++++++
 test/unit/BCoWHelper.t.sol           |  4 ++++
 test/unit/BCoWHelper.tree            |  1 +
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/contracts/BCoWHelper.sol b/src/contracts/BCoWHelper.sol
index 09a5b20f..4562bcda 100644
--- 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
index 54158906..003b4ce2 100644
--- 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
index 871b66f9..276cd68a 100644
--- 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
index ebbb7387..1cb8d994 100644
--- 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