test: mocks, conftest, twa tests (#9)

* merge-needed commits

* exports

* test: testsuite improvement, conftest & mocks, twa tests

* fix merging issues (double comment)

* addressing reviews

* chore: reuse interface from module

---------

Co-authored-by: Alberto <albicento.ac@gmail.com>

contracts/RewardsHandler.vy

@@ -34,32 +34,15 @@ implements: IERC165
 # yearn vault's interface
 from interfaces import IVault
 
-# crvUSD controller factory interface
-# used to compute the circulating supply
-from interfaces import IControllerFactory
-
 # we use access control because we want
 # to have multiple addresses being able to
 # adjust the rate while only the dao
 # (which has the `DEFAULT_ADMIN_ROLE`)
 # can appoint `RATE_MANAGER`s
 from snekmate.auth import access_control
 
-# import custom modules that contain
-# helper functions.
-import StablecoinLens as lens
-import TWA as twa
-
 initializes: access_control
-initializes: twa
-initializes: lens
-
 exports: (
-    # TODO add missing getters
-    twa.compute_twa,
-    twa.snapshots,
-    twa.get_len_snapshots,
-    twa.twa_window,
     # we don't expose `supportsInterface` from access control
     access_control.grantRole,
     access_control.revokeRole,
@@ -70,6 +53,22 @@ exports: (
     access_control.getRoleAdmin,
 )
 
+# import custom modules that contain
+# helper functions.
+import StablecoinLens as lens
+initializes: lens
+
+import TWA as twa
+initializes: twa
+exports: (
+    twa.compute_twa,
+    twa.snapshots,
+    twa.get_len_snapshots,
+    twa.twa_window,
+    twa.min_snapshot_dt_seconds,
+    twa.last_snapshot_timestamp,
+)
+
 RATE_MANAGER: public(constant(bytes32)) = keccak256("RATE_MANAGER")
 WEEK: constant(uint256) = 86400 * 7  # 7 days
 
@@ -98,7 +97,7 @@ def __init__(
     _stablecoin: IERC20,
     _vault: IVault,
     minimum_weight: uint256,
-    controller_factory: IControllerFactory,
+    controller_factory: lens.IControllerFactory,
     admin: address,
 ):
     lens.__init__(controller_factory)
@@ -150,7 +149,7 @@ def take_snapshot():
 
     supply_ratio: uint256 = supply_in_vault * 10**18 // circulating_supply
 
-    twa.store_snapshot(supply_ratio)
+    twa._store_snapshot(supply_ratio)
 
 
 @external
@@ -162,9 +161,7 @@ def process_rewards():
 
     # prevent the rewards from being distributed untill
     # the distribution rate has been set
-    assert (
-        self.distribution_time != 0
-    ), "rewards should be distributed over time"
+    assert (self.distribution_time != 0), "rewards should be distributed over time"
 
 
     # any crvUSD sent to this contract (usually
@@ -214,7 +211,7 @@ def weight() -> uint256:
         future if someone tries to manipulate the
         time-weighted average of the tvl ratio.
     """
-    return max(twa.compute(), self.minimum_weight)
+    return max(twa._compute(), self.minimum_weight)
 
 
 ################################################################
@@ -223,15 +220,15 @@ def weight() -> uint256:
 
 
 @external
-def set_twa_frequency(_min_snapshot_dt_seconds: uint256):
+def set_twa_snapshot_dt(_min_snapshot_dt_seconds: uint256):
     """
     @notice Setter for the time-weighted average minimal
         frequency.
     @param _min_snapshot_dt_seconds The minimum amount of
         time that should pass between two snapshots.
     """
     access_control._check_role(RATE_MANAGER, msg.sender)
-    twa.set_min_snapshot_dt_seconds(_min_snapshot_dt_seconds)
+    twa._set_snapshot_dt(_min_snapshot_dt_seconds)
 
 
 @external
@@ -242,7 +239,7 @@ def set_twa_window(_twa_window: uint256):
         the TWA value of the balance/supply ratio.
     """
     access_control._check_role(RATE_MANAGER, msg.sender)
-    twa.set_twa_window(_twa_window)
+    twa._set_twa_window(_twa_window)
 
 
 @external
@@ -291,6 +288,4 @@ def recover_erc20(token: IERC20, receiver: address):
     # to a trusted address.
     balance_to_recover: uint256 = staticcall token.balanceOf(self)
 
-    assert extcall token.transfer(
-        receiver, balance_to_recover, default_return_value=True
-    )
+    assert extcall token.transfer(receiver, balance_to_recover, default_return_value=True)

contracts/TWA.vy

@@ -9,7 +9,7 @@
 - TWA Calculation: Computes the TWA by iterating over the stored snapshots in reverse chronological order.
     It uses the trapezoidal rule to calculate the weighted average of the tracked value over the specified time window (`twa_window`).
 - Functions:
-  - `store_snapshot`: Internal function to store a new snapshot of the tracked value if the minimum time interval has passed.
+  - `_store_snapshot`: Internal function to store a new snapshot of the tracked value if the minimum time interval has passed.
         !!!Wrapper must be implemented in importing contract.
   - `compute_twa`: External view function that calculates and returns the TWA based on the stored snapshots.
   - `get_len_snapshots`: External view function that returns the total number of snapshots stored.
@@ -35,19 +35,6 @@ def __init__(_twa_window: uint256, _min_snapshot_dt_seconds: uint256):
     self.min_snapshot_dt_seconds = _min_snapshot_dt_seconds  # >=1s to prevent spamming
 
 
-@internal
-def store_snapshot(_value: uint256):
-    """
-    @notice Stores a snapshot of the tracked value.
-    @param _value The value to store.
-    """
-    if self.last_snapshot_timestamp + self.min_snapshot_dt_seconds <= block.timestamp:
-        self.last_snapshot_timestamp = block.timestamp
-        self.snapshots.append(
-            Snapshot(tracked_value=_value, timestamp=block.timestamp)
-        )  # store the snapshot into the DynArray
-
-
 @external
 @view
 def get_len_snapshots() -> uint256:
@@ -64,11 +51,24 @@ def compute_twa() -> uint256:
     """
     @notice External endpoint for _compute() function.
     """
-    return self.compute()
+    return self._compute()
+
+
+@internal
+def _store_snapshot(_value: uint256):
+    """
+    @notice Stores a snapshot of the tracked value.
+    @param _value The value to store.
+    """
+    if self.last_snapshot_timestamp + self.min_snapshot_dt_seconds <= block.timestamp:
+        self.last_snapshot_timestamp = block.timestamp
+        self.snapshots.append(
+            Snapshot(tracked_value=_value, timestamp=block.timestamp)
+        )  # store the snapshot into the DynArray
 
 
 @internal
-def set_twa_window(_new_window: uint256):
+def _set_twa_window(_new_window: uint256):
     """
     @notice Adjusts the TWA window.
     @param _new_window The new TWA window in seconds.
@@ -78,7 +78,7 @@ def set_twa_window(_new_window: uint256):
 
 
 @internal
-def set_min_snapshot_dt_seconds(_new_dt_seconds: uint256):
+def _set_snapshot_dt(_new_dt_seconds: uint256):
     """
     @notice Adjusts the minimum snapshot time interval.
     @param _new_dt_seconds The new minimum snapshot time interval in seconds.
@@ -89,7 +89,7 @@ def set_min_snapshot_dt_seconds(_new_dt_seconds: uint256):
 
 @internal
 @view
-def compute() -> uint256:
+def _compute() -> uint256:
     """
     @notice Computes the TWA over the specified time window by iterating backwards over the snapshots.
     @return The TWA for tracked value over the self.twa_window (10**18 decimals precision).

scripts/debug_tests.py

@@ -0,0 +1,29 @@
+import sys
+
+import pytest
+
+
+def is_debug_mode():
+    return sys.gettrace() is not None
+
+
+def main():
+    # Pytest arguments
+    pytest_args = [
+        "-s",  # Do not capture output, allowing you to see print statements and debug info
+        "tests/unitary/twa",  # Specific test to run
+        # '--maxfail=1',  # Stop after the firstD failure
+        "--tb=short",  # Shorter traceback for easier reading
+        "-rA",  # Show extra test summary info
+    ]
+
+    if not is_debug_mode():
+        pass
+        pytest_args.append("-n=auto")  # Automatically determine the number of workers
+
+    # Run pytest with the specified arguments
+    pytest.main(pytest_args)
+
+
+if __name__ == "__main__":
+    main()

tests/integration/conftest.py

@@ -13,7 +13,12 @@ def rpc_url():
 
 @pytest.fixture(scope="module", autouse=True)
 def forked_env(rpc_url):
-    boa.fork(rpc_url, block_identifier=18801970)
+    block_to_fork = 20742069
+    with boa.swap_env(boa.Env()):
+        boa.fork(url=rpc_url, block_identifier=block_to_fork)
+        print(f"\nForked the chain on block {boa.env.evm.vm.state.block_number}!")
+        boa.env.enable_fast_mode()
+        yield
 
 
 @pytest.fixture(scope="module")

tests/mocks/MockController.vy

@@ -8,6 +8,7 @@ implements: IController
 _monetary_policy: address
 total_debt: public(uint256)
 
+
 @external
 @view
 def monetary_policy() -> IMonetaryPolicy:

tests/mocks/MockControllerFactory.vy

@@ -7,11 +7,13 @@ implements: IControllerFactory
 
 _controllers: DynArray[IController, 10000]
 
+
 @external
 @view
 def controllers(i: uint256) -> IController:
     return self._controllers[i]
 
+
 @external
 @view
 def n_collaterals() -> uint256:

tests/mocks/MockMonetaryPolicy.vy

@@ -5,9 +5,10 @@ from contracts.interfaces import IPegKeeper
 
 implements: IMonetaryPolicy
 
-peg_keeper: IPegKeeper
+peg_keeper_array: IPegKeeper[1000]
+
 
 @external
 @view
 def peg_keepers(i: uint256) -> IPegKeeper:
-    return self.peg_keeper
+    return self.peg_keeper_array[i]

tests/mocks/MockPegKeeper.vy

@@ -0,0 +1,8 @@
+# pragma version ~=0.4
+
+debt: public(uint256)
+
+
+@deploy
+def __init__(circulating_supply: uint256):
+    self.debt = circulating_supply

tests/unitary/conftest.py

@@ -1,12 +1,24 @@
 import boa
 import pytest
 
+MOCK_CRV_USD_CIRCULATING_SUPPLY = 69_420_000 * 10**18
+
 
 @pytest.fixture(scope="module")
 def yearn_gov():
     return boa.env.generate_address()
 
 
+@pytest.fixture(scope="module")
+def curve_dao():
+    return boa.env.generate_address()
+
+
+@pytest.fixture(scope="module")
+def deployer():
+    return boa.env.generate_address()
+
+
 @pytest.fixture(scope="module")
 def vault_original():
     return boa.load("contracts/yearn/Vault.vy")
@@ -41,15 +53,6 @@ def vault(vault_factory, crvusd, role_manager):
     return vault_deployer.at(address)
 
 
-@pytest.fixture(scope="module")
-def lens():
-    lens = boa.load("tests/mocks/MockLens.vy")
-
-    lens.eval("self.supply = 1_000_000_000 * 10 ** 18")
-
-    return lens
-
-
 @pytest.fixture(scope="module")
 def vault_god(vault, role_manager):
     _god = boa.env.generate_address()
@@ -59,26 +62,53 @@ def vault_god(vault, role_manager):
     return _god
 
 
-@pytest.fixture(scope="module")
-def curve_dao():
-    return boa.env.generate_address()
-
-
 @pytest.fixture(params=[10**17, 5 * 10**17], scope="module")
 def minimum_weight(request):
     # TODO probably want to do some stateful testing here
     return request.param
 
 
 @pytest.fixture(scope="module")
-def controller_factory():
-    return boa.load("tests/mocks/MockControllerFactory.vy")
+def mock_controller_factory(mock_controller):
+    mock_controller_factory = boa.load("tests/mocks/MockControllerFactory.vy")
+    for i in range(4):  # because we use 3rd controller (weth) in contract code
+        mock_controller_factory.eval(
+            f"self._controllers.append(IController({mock_controller.address}))"
+        )
+    return mock_controller_factory
+
+
+@pytest.fixture(scope="module")
+def mock_controller(mock_monetary_policy):
+    mock_controller = boa.load("tests/mocks/MockController.vy")
+    mock_controller.eval(f"self._monetary_policy={mock_monetary_policy.address}")
+    return mock_controller
+
+
+@pytest.fixture(scope="module")
+def mock_monetary_policy(mock_peg_keeper):
+    mock_monetary_policy = boa.load("tests/mocks/MockMonetaryPolicy.vy")
+    mock_monetary_policy.eval(f"self.peg_keeper_array[0] = IPegKeeper({mock_peg_keeper.address})")
+    return mock_monetary_policy
+
+
+@pytest.fixture(scope="module")
+def mock_peg_keeper():
+    mock_peg_keeper = boa.load("tests/mocks/MockPegKeeper.vy", MOCK_CRV_USD_CIRCULATING_SUPPLY)
+    return mock_peg_keeper
 
 
 @pytest.fixture(scope="module")
-def rewards_handler(vault, crvusd, role_manager, minimum_weight, controller_factory, curve_dao):
+def rewards_handler(
+    vault, crvusd, role_manager, minimum_weight, mock_controller_factory, curve_dao
+):
     rh = boa.load(
-        "contracts/RewardsHandler.vy", crvusd, vault, minimum_weight, controller_factory, curve_dao
+        "contracts/RewardsHandler.vy",
+        crvusd,
+        vault,
+        minimum_weight,
+        mock_controller_factory,
+        curve_dao,
     )
 
     vault.set_role(rh, 2**11 | 2**5 | 2**0, sender=role_manager)