diff --git a/contracts/TornadoTrees.sol b/contracts/TornadoTrees.sol
index d550513..6e23475 100644
--- a/contracts/TornadoTrees.sol
+++ b/contracts/TornadoTrees.sol
@@ -235,10 +235,7 @@ contract TornadoTrees is Initializable {
     string memory _type,
     uint256 _from, // most likely array length after the proposal has passed
     uint256 _step // optimal step size to find first match, approximately equals dispersion
-  ) internal view returns (uint256) {
-    if (_from == 0 && _step == 0) {
-      return 0; // for tests
-    }
+  ) internal view virtual returns (uint256) {
     // Find the segment with correct array length
     bool direction = elementExists(_tornadoTreesV1, _type, _from);
     do {
diff --git a/contracts/mocks/TornadoTreesMock.sol b/contracts/mocks/TornadoTreesMock.sol
index cb65be3..adcd073 100644
--- a/contracts/mocks/TornadoTreesMock.sol
+++ b/contracts/mocks/TornadoTreesMock.sol
@@ -113,4 +113,16 @@ contract TornadoTreesMock is TornadoTrees {
       _withdrawals[i] = withdrawals[lastProcessedWithdrawalLeaf + i];
     }
   }
+
+  function findArrayLength(
+    ITornadoTreesV1 _tornadoTreesV1,
+    string memory _type,
+    uint256 _from, // most likely array length after the proposal has passed
+    uint256 _step // optimal step size to find first match, approximately equals dispersion
+  ) internal view override returns (uint256) {
+    if (_from == 0 && _step == 0) {
+      return 0;
+    }
+    return super.findArrayLength(_tornadoTreesV1, _type, _from, _step);
+  }
 }