Update integration tests to handle hook data

python/src/hooks/exit_fee/types.py

@@ -6,3 +6,35 @@ class ExitFeeHookState:
     remove_liquidity_hook_fee_percentage: int
     tokens: list[str]
     hook_type: str = "ExitFee"
+
+
+def map_exit_fee_hook_state(hook_data: dict, tokens: list[str]) -> ExitFeeHookState:
+    """
+    Maps EXIT_FEE hook data to ExitFeeHookState.
+
+    Args:
+        hook_data: Raw hook dict from JSON with dynamicData containing
+                   removeLiquidityHookFeePercentage
+        tokens: List of token addresses from pool data
+
+    Returns:
+        ExitFeeHookState object
+
+    Raises:
+        ValueError: If required fields are missing
+    """
+    if "dynamicData" not in hook_data:
+        raise ValueError("EXIT_FEE hook requires dynamicData")
+
+    dynamic_data = hook_data["dynamicData"]
+    if "removeLiquidityHookFeePercentage" not in dynamic_data:
+        raise ValueError(
+            "EXIT_FEE hook requires removeLiquidityHookFeePercentage in dynamicData"
+        )
+
+    return ExitFeeHookState(
+        remove_liquidity_hook_fee_percentage=int(
+            dynamic_data["removeLiquidityHookFeePercentage"]
+        ),
+        tokens=tokens,
+    )

python/src/pools/stable/stable.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from src.common.maths import Rounding, mul_down_fixed
+from src.common.maths import Rounding, mul_up_fixed
 from src.common.pool_base import PoolBase
 from src.common.swap_params import SwapParams
 from src.common.types import SwapKind
@@ -67,7 +67,7 @@ def compute_balance(
         return compute_balance(
             self.amp,
             balances_live_scaled18,
-            mul_down_fixed(
+            mul_up_fixed(
                 self.compute_invariant(balances_live_scaled18, Rounding.ROUND_UP),
                 invariant_ratio,
             ),

python/src/vault/add_liquidity.py

@@ -15,6 +15,7 @@
     _get_single_input_index,
     _require_unbalanced_liquidity_enabled,
     _to_raw_undo_rate_round_up,
+    _to_scaled_18_apply_rate_round_down,
 )
 from src.hooks.types import HookBase, HookState
 
@@ -103,14 +104,20 @@ def add_liquidity(
 
         # A Pool's token balance always decreases after an exit
         # Computes protocol and pool creator fee which is eventually taken from pool balance
-        aggregate_swap_fee_amount_scaled18 = _compute_and_charge_aggregate_swap_fees(
+        aggregate_swap_fee_amount_raw = _compute_and_charge_aggregate_swap_fees(
             swap_fee_amounts_scaled18[i],
             pool_state.aggregate_swap_fee,
             pool_state.scaling_factors,
             pool_state.token_rates,
             i,
         )
 
+        aggregate_swap_fee_amount_scaled18 = _to_scaled_18_apply_rate_round_down(
+            aggregate_swap_fee_amount_raw,
+            pool_state.scaling_factors[i],
+            pool_state.token_rates[i],
+        )
+
         # Update the balances with the incoming amounts and subtract the swap fees
         updated_balances_live_scaled18[i] = (
             updated_balances_live_scaled18[i]

python/src/vault/swap.py

@@ -206,10 +206,11 @@ def _compute_amount_given_scaled18(
             token_rates[index_in],
         )
     else:
+        rate_rounded_up = _compute_rate_round_up(token_rates[index_out])
         amount_given_scaled_18 = _to_scaled_18_apply_rate_round_up(
             amount_given_raw,
             scaling_factors[index_out],
-            token_rates[index_out],
+            rate_rounded_up,
         )
 
     return amount_given_scaled_18

python/test/test_add_liquidity.py

@@ -1,4 +1,7 @@
-from test.utils.map_pool_state import map_pool_state, transform_strings_to_ints
+from test.utils.map_pool_state import (
+    map_pool_and_hook_state,
+    transform_strings_to_ints,
+)
 from test.utils.read_test_data import read_test_data
 from typing import cast
 
@@ -19,16 +22,23 @@ def test_add_liquidity():
             raise ValueError("Buffer pools do not support addLiquidity")
         # note any amounts must be passed as ints not strings
         pool_with_ints = transform_strings_to_ints(pool)
+        pool_state, hook_state = map_pool_and_hook_state(pool_with_ints)
         calculated_amount = vault.add_liquidity(
             add_liquidity_input=AddLiquidityInput(
                 pool=pool["poolAddress"],
                 max_amounts_in_raw=list(map(int, add_test["inputAmountsRaw"])),
                 min_bpt_amount_out_raw=int(add_test["bptOutRaw"]),
                 kind=AddLiquidityKind(add_test["kind"]),
             ),
-            pool_state=cast(PoolState, map_pool_state(pool_with_ints)),
+            pool_state=cast(PoolState, pool_state),
+            hook_state=hook_state,
         )
-        assert calculated_amount.bpt_amount_out_raw == int(add_test["bptOutRaw"])
+        # Relax test assertion to accept off-by-1 error because testData might
+        # return amounts off-by-1 when compared to actual implementations.
+        # e.g. getCurrentLiveBalances rounds pools balances down, while solidity
+        # rounds pool balances up when loading pool data within add liquidity operations
+        assert calculated_amount.bpt_amount_out_raw >= int(add_test["bptOutRaw"]) - 1
+        assert calculated_amount.bpt_amount_out_raw <= int(add_test["bptOutRaw"]) + 1
         assert calculated_amount.amounts_in_raw == list(
             map(int, add_test["inputAmountsRaw"])
         )

python/test/test_remove_liquidity.py

@@ -1,4 +1,7 @@
-from test.utils.map_pool_state import map_pool_state, transform_strings_to_ints
+from test.utils.map_pool_state import (
+    map_pool_and_hook_state,
+    transform_strings_to_ints,
+)
 from test.utils.read_test_data import read_test_data
 from typing import cast
 
@@ -19,14 +22,16 @@ def test_remove_liquidity():
             raise ValueError("Buffer pools do not support addLiquidity")
         # note any amounts must be passed as ints not strings
         pool_with_ints = transform_strings_to_ints(pool)
+        pool_state, hook_state = map_pool_and_hook_state(pool_with_ints)
         calculated_amount = vault.remove_liquidity(
             remove_liquidity_input=RemoveLiquidityInput(
                 pool=pool["poolAddress"],
                 min_amounts_out_raw=list(map(int, remove_test["amountsOutRaw"])),
                 max_bpt_amount_in_raw=int(remove_test["bptInRaw"]),
                 kind=RemoveLiquidityKind(remove_test["kind"]),
             ),
-            pool_state=cast(PoolState, map_pool_state(pool_with_ints)),
+            pool_state=cast(PoolState, pool_state),
+            hook_state=hook_state,
         )
         assert calculated_amount.bpt_amount_in_raw == int(remove_test["bptInRaw"])
         assert calculated_amount.amounts_out_raw == list(

python/test/test_swaps.py

@@ -1,4 +1,7 @@
-from test.utils.map_pool_state import map_pool_state, transform_strings_to_ints
+from test.utils.map_pool_state import (
+    map_pool_and_hook_state,
+    transform_strings_to_ints,
+)
 from test.utils.read_test_data import read_test_data
 
 from src.common.types import SwapInput, SwapKind
@@ -18,14 +21,16 @@ def test_swaps():
         pool = test_data["pools"][swap_test["test"]]
         # note any amounts must be passed as ints not strings
         pool_with_ints = transform_strings_to_ints(pool)
+        pool_state, hook_state = map_pool_and_hook_state(pool_with_ints)
         calculated_amount = vault.swap(
             swap_input=SwapInput(
                 amount_raw=int(swap_test["amountRaw"]),
                 token_in=swap_test["tokenIn"],
                 token_out=swap_test["tokenOut"],
                 swap_kind=SwapKind(swap_test["swapKind"]),
             ),
-            pool_state=map_pool_state(pool_with_ints),
+            pool_state=pool_state,
+            hook_state=hook_state,
         )
         if pool["poolType"] == "Buffer":
             assert are_big_ints_within_percent(