test: testsuite improvement, conftest & mocks, twa tests

contracts/RewardsHandler.vy

@@ -45,11 +45,6 @@ from interfaces import IControllerFactory
 # 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
 exports: (
     # we don't expose `supportsInterface` from access control
@@ -62,8 +57,12 @@ 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,
@@ -120,6 +119,7 @@ def __init__(
     stablecoin = _stablecoin
     vault = _vault
 
+
 ################################################################
 #                   PERMISSIONLESS FUNCTIONS                   #
 ################################################################
@@ -159,10 +159,6 @@ def take_snapshot():
 
     twa._store_snapshot(supply_ratio)
 
-    # get the circulating supply from a helper function
-    # (supply in circulation = controllers' debt + peg
-    # keppers' debt)
-    circulating_supply: uint256 = lens._circulating_supply()
 
 @external
 def process_rewards():
@@ -173,9 +169,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
@@ -302,6 +296,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

@@ -67,7 +67,6 @@ def _store_snapshot(_value: uint256):
         )  # store the snapshot into the DynArray
 
 
-
 @internal
 def _set_twa_window(_new_window: uint256):
     """

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,6 +1,8 @@
 import boa
 import pytest
 
+MOCK_CRV_USD_CIRCULATING_SUPPLY = 69_420_000 * 10**18
+
 
 @pytest.fixture(scope="module")
 def yearn_gov():
@@ -60,28 +62,53 @@ def vault_god(vault, role_manager):
     return _god
 
 
-@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(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)

tests/unitary/twa/test_internal_compute.py

@@ -1,15 +1,12 @@
 import boa
 
 
-def test_compute_twa_no_snapshots(rewards_handler, crvusd, vault, lens, vault_god):
+def test_compute_twa_no_snapshots(rewards_handler, crvusd, vault, vault_god):
     # Prepare Alice's balance
     alice = boa.env.generate_address()
     boa.deal(crvusd, alice, 100_000_000 * 10**18)
     crvusd.approve(vault.address, 2**256 - 1, sender=alice)
 
-    # Set up RewardsHandler
-    rewards_handler.eval(f"self.lens = {lens.address}")
-
     # Set vault deposit limit
     vault.set_deposit_limit(crvusd.balanceOf(alice), sender=vault_god)
 
@@ -29,28 +26,26 @@ def test_one_deposit_boa(rewards_handler):
     assert rewards_handler.compute_twa() == 100 * 10**18
 
 
-def test_many_deposit_boa(rewards_handler):
+def test_many_deposits_boa(rewards_handler):
     TWA_WINDOW = 1000
     rewards_handler.eval(f"twa.twa_window = {TWA_WINDOW}")
     N_ITER = 10
     AMT_ADD = 10 * 10**18
     for i in range(N_ITER):
-        rewards_handler.eval(f"twa._store_snapshot({AMT_ADD})")
+        snapshot_amount = AMT_ADD * (i + 1)
+        rewards_handler.eval(f"twa._store_snapshot({snapshot_amount})")
         boa.env.time_travel(seconds=TWA_WINDOW // N_ITER)
-    assert rewards_handler.compute_twa() < AMT_ADD * N_ITER
-    boa.env.time_travel(seconds=2000)
-    assert rewards_handler.compute_twa() == 100 * 10**18
+    assert rewards_handler.compute_twa() < snapshot_amount
+    boa.env.time_travel(seconds=TWA_WINDOW)
+    assert rewards_handler.compute_twa() == snapshot_amount
 
 
-def test_twa_one_deposit(vault, crvusd, rewards_handler, lens, vault_god):
+def test_twa_one_deposit(vault, crvusd, rewards_handler, vault_god):
     # Prepare Alice's balance
     alice = boa.env.generate_address()
     boa.deal(crvusd, alice, 100_000_000 * 10**18)
     crvusd.approve(vault.address, 2**256 - 1, sender=alice)
 
-    # Set up RewardsHandler
-    rewards_handler.eval(f"self.lens = {lens.address}")
-
     # Set vault deposit limit
     vault.set_deposit_limit(crvusd.balanceOf(alice), sender=vault_god)
 
@@ -64,19 +59,19 @@ def test_twa_one_deposit(vault, crvusd, rewards_handler, lens, vault_god):
 
     # Compute TWA
     twa = rewards_handler.compute_twa()
-    expected_twa = AMT_DEPOSIT * 10**18 // lens.circulating_supply()
+
+    circulating_supply = rewards_handler.eval("lens._circulating_supply()")
+    expected_twa = AMT_DEPOSIT * 10**18 // circulating_supply
+
     assert twa == expected_twa, "TWA does not match expected amount"
 
 
-def test_twa_trapezoid(vault, crvusd, rewards_handler, lens, vault_god):
+def test_twa_trapezoid(vault, crvusd, rewards_handler, vault_god):
     # Prepare Alice's balance
     alice = boa.env.generate_address()
     boa.deal(crvusd, alice, 100_000_000 * 10**18)
     crvusd.approve(vault.address, 2**256 - 1, sender=alice)
 
-    # Set up RewardsHandler
-    rewards_handler.eval(f"self.lens = {lens.address}")
-
     # Set vault deposit limit
     vault.set_deposit_limit(crvusd.balanceOf(alice), sender=vault_god)
 
@@ -95,19 +90,17 @@ def test_twa_trapezoid(vault, crvusd, rewards_handler, lens, vault_god):
     # Compute TWA
     twa = rewards_handler.compute_twa()
     # 1.5 * AMT_DEPOSIT because we have two equal deposits and trapezoidal rule
-    expected_twa = 1.5 * AMT_DEPOSIT * 10**18 // lens.circulating_supply()
+    circulating_supply = rewards_handler.eval("lens._circulating_supply()")
+    expected_twa = 1.5 * AMT_DEPOSIT * 10**18 // circulating_supply
     assert twa == expected_twa, "TWA does not match expected amount"
 
 
-def test_twa_multiple_deposits(vault, crvusd, rewards_handler, lens, vault_god):
+def test_twa_multiple_deposits(vault, crvusd, rewards_handler, vault_god):
     # Prepare Alice's balance
     alice = boa.env.generate_address()
     boa.deal(crvusd, alice, 100_000_000 * 10**18)
     crvusd.approve(vault.address, 2**256 - 1, sender=alice)
 
-    # Set up RewardsHandler
-    rewards_handler.eval(f"self.lens = {lens.address}")
-
     # Set vault deposit limit to Alice's full balance
     vault.set_deposit_limit(crvusd.balanceOf(alice), sender=vault_god)
 
@@ -186,10 +179,10 @@ def test_twa_multiple_deposits(vault, crvusd, rewards_handler, lens, vault_god):
 
     # Calculate the actual staked rate for comparison
     total_staked_amount = AMT_DEPOSIT * N_ITERATIONS
-    circulating_supply = lens.circulating_supply()
+    circulating_supply = rewards_handler.eval("lens._circulating_supply()")
     staked_rate = total_staked_amount * 10**18 // circulating_supply
 
-    print(f"Staked rate: {staked_rate}, Contract TWA: {twa}")
+    # print(f"Staked rate: {staked_rate}, Contract TWA: {twa}")
 
     # Compare the TWA from the contract against the expected values
     assert twa <= staked_rate, "TWA is unexpectedly higher than the staked rate"