Stateful testing of TWA module

.github/workflows/ci.yaml

@@ -3,23 +3,21 @@ name: CI Workflow
 on: [push]
 
 jobs:
-
   # Linting with Pre-commit
-  lint:
+  ruff:
     runs-on: ubuntu-latest
-
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v1
 
-      - name: Setup Python 3.12.6
+      - name: Set up Python 3.12.6
         uses: actions/setup-python@v5
         with:
           python-version: 3.12.6
 
       # Run pre-commit hooks (e.g., black, flake8, isort)
       - uses: pre-commit/[email protected]
 
-
   # Test runs
   tests:
     runs-on: ubuntu-latest
@@ -28,31 +26,24 @@ jobs:
       fail-fast: false
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
-      # Install Solidity compiler (solc)
-      - name: Install solc
-        run: |
-          sudo add-apt-repository -y ppa:ethereum/ethereum
-          sudo apt-get update
-          sudo apt-get install -y solc
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "0.4.18"
+          enable-cache: true  # Enables built-in caching for uv
 
-      # install poetry to allow cacheing
-      - name: Install poetry
-        run: pipx install poetry
+      - name: Set up Python 3.12.6
+        run: uv python install 3.12.6
 
-      - name: Setup Python 3.12.6
-        uses: actions/setup-python@v5
-        with:
-          python-version: 3.12.6
-          cache: 'poetry'
+      # Install dependencies with all extras (including dev)
       - name: Install Requirements
-        run: |
-          poetry config virtualenvs.in-project true
-          poetry install
+        run: uv sync --extra=dev
+
+      # Run tests with environment variables
       - name: Run Tests
         env:
           ETH_RPC_URL: ${{ secrets.ETH_RPC_URL }}
           ETHERSCAN_API_KEY: ${{ secrets.ETHERSCAN_API_KEY }}
-        run: |
-          poetry run pytest -n auto
+        run: uv run pytest -n auto

.pre-commit-config.yaml

@@ -1,6 +1,8 @@
+exclude: 'contracts/yearn/VaultV3\.vy|VaultFactory\.vy'
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -9,20 +11,12 @@ repos:
       - id: check-ast
       - id: detect-private-key
 
-  - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 24.8.0
-    hooks:
-      - id: black
-        args: [--line-length=100]
-  - repo: https://github.com/pycqa/flake8
-    rev: 7.1.1
-    hooks:
-      - id: flake8
-  - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.6.9
     hooks:
-      - id: isort
-        args: ["--profile", "black", --line-length=100]
+      - id: ruff
+        args: [--fix, --exit-non-zero-on-fix]
+      - id: ruff-format
 
 default_language_version:
   python: python3.12

contracts/DepositLimitModule.vy

@@ -0,0 +1,167 @@
+# pragma version ~=0.4
+
+"""
+@title Temporary Deposit Limit Module for yearn vaults
+@notice This contract temporarily controls deposits into a yearn vault and manages roles (admin and security_agent).
+@dev Admins can appoint new admins and security_agents, while security_agents can pause or resume deposits.
+This module will be removed once the vault is battle-tested and proven stable.
+"""
+
+
+################################################################
+#                           INTERFACES                         #
+################################################################
+
+from interfaces import IVault
+
+################################################################
+#                           STORAGE                            #
+################################################################
+
+# Two main roles: admin and security_agent
+is_admin: public(HashMap[address, bool])
+is_security_agent: public(HashMap[address, bool])
+
+# Main contract state for deposit control
+deposits_paused: public(bool)
+
+# Max deposit limit
+max_deposit_limit: public(uint256)
+
+# Stablecoin/Vault addresses
+vault: public(immutable(IVault))
+
+################################################################
+#                         CONSTRUCTOR                          #
+################################################################
+
+@deploy
+def __init__(
+    _vault: IVault,
+    max_deposit_limit: uint256,
+):
+    """
+    @notice Initializes the contract by assigning the deployer as the initial admin and security_agent.
+    """
+    self._set_admin(msg.sender, True)
+    self._set_security_agent(msg.sender, True)
+    self._set_deposits_paused(False)  # explicit non-paused at init
+    self._set_deposit_limit(max_deposit_limit)
+
+    vault = _vault
+
+
+################################################################
+#                      INTERNAL FUNCTIONS                      #
+################################################################
+
+@internal
+def _set_admin(_address: address, is_admin: bool):
+    """
+    @notice Internal function to assign or revoke admin role.
+    @param address The address to be granted or revoked admin status.
+    @param is_admin Boolean indicating if the address should be an admin.
+    """
+    self.is_admin[_address] = is_admin
+
+
+@internal
+def _set_security_agent(_address: address, is_security_agent: bool):
+    """
+    @notice Internal function to assign or revoke security_agent role.
+    @param address The address to be granted or revoked security_agent status.
+    @param is_security_agent Boolean indicating if the address should be a security_agent.
+    """
+    self.is_security_agent[_address] = is_security_agent
+
+
+@internal
+def _set_deposits_paused(is_paused: bool):
+    """
+    @notice Internal function to pause or unpause deposits.
+    @param is_paused Boolean indicating if deposits should be paused.
+    """
+    self.deposits_paused = is_paused
+
+
+@internal
+def _set_deposit_limit(new_limit: uint256):
+    """
+    @notice Internal function to set the maximum deposit limit.
+    @param new_limit The new maximum deposit limit.
+    """
+    self.max_deposit_limit = new_limit
+
+
+################################################################
+#                        EXTERNAL FUNCTIONS                    #
+################################################################
+
+@external
+def set_admin(new_admin: address, is_admin: bool):
+    """
+    @notice Allows an admin to grant or revoke admin role to another address.
+    @param new_admin The address to grant or revoke admin role.
+    @param is_admin Boolean indicating if the address should be an admin.
+    @dev Only callable by an admin.
+    """
+    assert self.is_admin[msg.sender], "Caller is not an admin"
+    self._set_admin(new_admin, is_admin)
+
+
+@external
+def set_security_agent(new_security_agent: address, is_security_agent: bool):
+    """
+    @notice Allows an admin to grant or revoke security_agent role to another address.
+    @param new_security_agent The address to grant or revoke security_agent role.
+    @param is_security_agent Boolean indicating if the address should be a security_agent.
+    @dev Only callable by an admin.
+    """
+    assert self.is_admin[msg.sender], "Caller is not an admin"
+    self._set_security_agent(new_security_agent, is_security_agent)
+
+
+@external
+def set_deposits_paused(state: bool):
+    """
+    @notice Allows a security_agent to pause or resume deposits.
+    @param state Boolean indicating the desired paused state for deposits.
+    @dev Only callable by a security_agent.
+    """
+    assert self.is_security_agent[msg.sender], "Caller is not a security_agent"
+    self._set_deposits_paused(state)
+
+
+@external
+def set_deposit_limit(new_limit: uint256):
+    """
+    @notice Allows an admin to update the maximum deposit limit.
+    @param new_limit The new maximum deposit limit.
+    @dev Only callable by an admin.
+    """
+    assert self.is_admin[msg.sender], "Caller is not an admin"
+    self._set_deposit_limit(new_limit)
+
+
+################################################################
+#                        VIEW FUNCTIONS                        #
+################################################################
+
+@view
+@external
+def available_deposit_limit(receiver: address) -> uint256:
+    """
+    @notice Checks the available deposit limit for a given receiver.
+    @param receiver The address querying deposit limit.
+    @return uint256 Returns the maximum deposit limit if deposits are not paused, otherwise returns 0.
+    """
+    if self.deposits_paused:
+        return 0
+    if self.max_deposit_limit == max_value(uint256):
+        return max_value(uint256)
+    else:
+        vault_balance: uint256 = staticcall vault.totalAssets()
+        if vault_balance >= self.max_deposit_limit:
+            return 0
+        else:
+            return self.max_deposit_limit - vault_balance

contracts/RewardsHandler.vy

@@ -83,6 +83,7 @@ exports: (
 event MinimumWeightUpdated:
     new_minimum_weight: uint256
 
+
 event ScalingFactorUpdated:
     new_scaling_factor: uint256
 
@@ -93,7 +94,8 @@ event ScalingFactorUpdated:
 
 
 RATE_MANAGER: public(constant(bytes32)) = keccak256("RATE_MANAGER")
-WEEK: constant(uint256) = 86400 * 7  # 7 days
+RECOVERY_MANAGER: public(constant(bytes32)) = keccak256("RECOVERY_MANAGER")
+WEEK: constant(uint256) = 86_400 * 7  # 7 days
 MAX_BPS: constant(uint256) = 10**4  # 100%
 
 _SUPPORTED_INTERFACES: constant(bytes4[3]) = [
@@ -142,8 +144,9 @@ def __init__(
     access_control.__init__()
     # admin (most likely the dao) controls who can be a rate manager
     access_control._grant_role(access_control.DEFAULT_ADMIN_ROLE, admin)
-    # admin itself is a RATE_ADMIN
+    # admin itself is a RATE_MANAGER and RECOVERY_MANAGER
     access_control._grant_role(RATE_MANAGER, admin)
+    access_control._grant_role(RECOVERY_MANAGER, admin)
     # deployer does not control this contract
     access_control._revoke_role(access_control.DEFAULT_ADMIN_ROLE, msg.sender)
 
@@ -204,12 +207,13 @@ def process_rewards():
         self.distribution_time != 0
     ), "rewards should be distributed over time"
 
-
     # any crvUSD sent to this contract (usually through the fee splitter, but
     # could also come from other sources) will be used as a reward for crvUSD
     # stakers in the vault.
     available_balance: uint256 = staticcall stablecoin.balanceOf(self)
 
+    assert available_balance > 0, "no rewards to distribute"
+
     # we distribute funds in 2 steps:
     # 1. transfer the actual funds
     extcall stablecoin.transfer(vault.address, available_balance)
@@ -247,7 +251,9 @@ def weight() -> uint256:
     for more at the beginning and can also be increased in the future if someone
     tries to manipulate the time-weighted average of the tvl ratio.
     """
-    return max(twa._compute() * self.scaling_factor // MAX_BPS, self.minimum_weight)
+    return max(
+        twa._compute() * self.scaling_factor // MAX_BPS, self.minimum_weight
+    )
 
 
 ################################################################
@@ -298,7 +304,8 @@ def set_distribution_time(new_distribution_time: uint256):
     extcall vault.setProfitMaxUnlockTime(new_distribution_time)
 
     # enact the changes
-    extcall vault.process_report(staticcall vault.default_queue(0))
+    extcall vault.process_report(vault.address)
+
 
 @external
 def set_minimum_weight(new_minimum_weight: uint256):
@@ -345,7 +352,7 @@ def recover_erc20(token: IERC20, receiver: address):
     to this contract. crvUSD cannot be recovered as it's part of the core logic of
     this contract.
     """
-    access_control._check_role(RATE_MANAGER, msg.sender)
+    access_control._check_role(RECOVERY_MANAGER, msg.sender)
 
     # if crvUSD was sent by accident to the contract the funds are lost and will
     # be distributed as staking rewards on the next `process_rewards` call.

contracts/TWA.vy

@@ -70,7 +70,7 @@ struct Snapshot:
 @deploy
 def __init__(_twa_window: uint256, _min_snapshot_dt_seconds: uint256):
     self._set_twa_window(_twa_window)
-    self._set_snapshot_dt(_min_snapshot_dt_seconds)
+    self._set_snapshot_dt(max(1, _min_snapshot_dt_seconds))
 
 
 ################################################################
@@ -158,15 +158,20 @@ def _compute() -> uint256:
         i_backwards: uint256 = index_array_end - i
         current_snapshot: Snapshot = self.snapshots[i_backwards]
         next_snapshot: Snapshot = current_snapshot
-        if i != 0:  # If not the first iteration, get the next snapshot
+        if i != 0:  # If not the first iteration (last snapshot), get the next snapshot
             next_snapshot = self.snapshots[i_backwards + 1]
 
+        # Time Axis (Increasing to the Right) --->
+        #                                        SNAPSHOT
+        # |---------|---------|---------|------------------------|---------|---------|
+        # t0   time_window_start      interval_start        interval_end      block.timestamp (Now)
+
         interval_start: uint256 = current_snapshot.timestamp
         # Adjust interval start if it is before the time window start
         if interval_start < time_window_start:
             interval_start = time_window_start
 
-        interval_end: uint256 = 0
+        interval_end: uint256 = interval_start
         if i == 0:  # First iteration - we are on the last snapshot (i_backwards = num_snapshots - 1)
             # For the last snapshot, interval end is block.timestamp
             interval_end = block.timestamp
@@ -186,7 +191,10 @@ def _compute() -> uint256:
         total_weighted_tracked_value += averaged_tracked_value * time_delta
         total_time += time_delta
 
+    if total_time == 0 and len(self.snapshots) == 1:
+        # case when only snapshot is taken in the block where computation is called
+        return self.snapshots[0].tracked_value
+
     assert total_time > 0, "Zero total time!"
     twa: uint256 = total_weighted_tracked_value // total_time
-
     return twa