Skip to content

Commit

Permalink
test(Python): Hooks - before swap hook.
Browse files Browse the repository at this point in the history
  • Loading branch information
johngrantuk committed Aug 19, 2024
1 parent 7471c96 commit 8db27a3
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 4 deletions.
6 changes: 4 additions & 2 deletions python/src/swap.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,12 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
# will mean something has gone wrong.
# We do take into account and balance changes due
# to hook using hookAdjustedBalancesScaled18.
hook_return = hook_class.onBeforeSwap({**swap_input, "hook_state": hook_state})
hook_return = hook_class.on_before_swap(
{**swap_input, "hook_state": hook_state}
)
if hook_return["success"] is False:
raise SystemError("BeforeSwapHookFailed")
for i, a in enumerate(hook_return["hookAdjustedBalancesScaled18"]):
for i, a in enumerate(hook_return["hook_adjusted_balances_scaled18"]):
updated_balances_live_scaled18[i] = a

swap_fee = pool_state["swapFee"]
Expand Down
4 changes: 2 additions & 2 deletions python/test/hooks/after_swap.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ def on_compute_dynamic_swap_fee(self):
custom_hook_classes={"CustomHook": CustomHook},
)

def test_hook_after_remove_liquidity_no_fee():
def test_hook_after_swap_no_fee():
# aggregateSwapFee of 0 should not take any protocol fees from updated balances
# hook state is used to pass expected value to tests
# with aggregateFee = 0, balance out is just balance - calculated
Expand All @@ -142,7 +142,7 @@ def test_hook_after_remove_liquidity_no_fee():
assert test == 1


def test_hook_after_add_liquidity_with_fee():
def test_hook_after_swap_with_fee():
# aggregateSwapFee of 50% should take half of remaining
# hook state is used to pass expected value to tests
# Aggregate fee amount is 50% of swap fee
Expand Down
118 changes: 118 additions & 0 deletions python/test/hooks/before_swap.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
import pytest
import sys
import os

# Get the directory of the current file
current_file_dir = os.path.dirname(os.path.abspath(__file__))
# Get the parent directory (one level up)
parent_dir = os.path.dirname(os.path.dirname(current_file_dir))

# Insert the parent directory at the start of sys.path
sys.path.insert(0, parent_dir)

from src.pools.weighted import Weighted
from src.swap import SwapKind

# Get the directory of the current file
current_file_dir = os.path.dirname(os.path.abspath(__file__))
# Get the parent directory (one level up)
parent_dir = os.path.dirname(os.path.dirname(current_file_dir))

# Insert the parent directory at the start of sys.path
sys.path.insert(0, parent_dir)

from src.vault import Vault
from src.hooks.default_hook import DefaultHook

pool = {
"poolType": "CustomPool",
"hookType": "CustomHook",
"chainId": "11155111",
"blockNumber": "5955145",
"poolAddress": "0xb2456a6f51530053bc41b0ee700fe6a2c37282e8",
"tokens": [
"0x7b79995e5f793A07Bc00c21412e50Ecae098E7f9",
"0xb19382073c7A0aDdbb56Ac6AF1808Fa49e377B75",
],
"scalingFactors": [1000000000000000000, 1000000000000000000],
"weights": [500000000000000000, 500000000000000000],
"swapFee": 100000000000000000,
"balancesLiveScaled18": [2000000000000000000, 2000000000000000000],
"tokenRates": [1000000000000000000, 1000000000000000000],
"totalSupply": 1000000000000000000,
"aggregateSwapFee": 500000000000000000,
}


swap_input = {
"amount_raw": 100000000,
"swap_kind": SwapKind.GIVENIN.value,
"token_in": pool['tokens'][0],
"token_out": pool['tokens'][1],
}

class CustomPool(Weighted):
def __init__(self, pool_state):
super().__init__(pool_state)

class CustomHook:
def __init__(self):
self.should_call_compute_dynamic_swap_fee = False
self.should_call_before_swap = True
self.should_call_after_swap = False
self.should_call_before_add_liquidity = False
self.should_call_after_add_liquidity = False
self.should_call_before_remove_liquidity = False
self.should_call_after_remove_liquidity = False
self.enable_hook_adjusted_amounts = False

def on_before_add_liquidity(self):
return {'success': False, 'hook_adjusted_balances_scaled18': []}

def on_after_add_liquidity(self, kind, amounts_in_scaled18, amounts_in_raw, bpt_amount_out, balances_scaled18, hook_state):
return { 'success': False, 'hook_adjusted_amounts_in_raw': [] };

def on_before_remove_liquidity(self):
return {'success': False, 'hook_adjusted_balances_scaled18': []}

def on_after_remove_liquidity(self, kind, bpt_amount_in, amounts_out_scaled18, amounts_out_raw, balances_scaled18, hook_state):
return {
'success': False,
'hook_adjusted_amounts_out_raw': []
}

def on_before_swap(self, params):
hook_state = params['hook_state']
if not (isinstance(hook_state, dict) and hook_state is not None and 'balanceChange' in hook_state):
raise ValueError('Unexpected hookState')
assert params['swap_kind'] == swap_input['swap_kind']
assert params['token_in'] == swap_input['token_in']
assert params['token_out'] == swap_input['token_out']
assert params['amount_raw'] == swap_input['amount_raw']
return {'success': True, 'hook_adjusted_balances_scaled18': hook_state['balanceChange']}

def on_after_swap(self, params):
return {'success': True, 'hook_adjusted_amount_calculated_raw': 0}

def on_compute_dynamic_swap_fee(self):
return {'success': False, 'dynamic_swap_fee': 0}

vault = Vault(
custom_pool_classes={"CustomPool": CustomPool},
custom_hook_classes={"CustomHook": CustomHook},
)

def test_before_swap():
# should alter pool balances
# hook state is used to pass new balances which give expected swap result
input_hook_state = {
"balanceChange": [
1000000000000000000, 1000000000000000000
],
}
test = vault.swap(
swap_input,
pool,
hook_state=input_hook_state
)
assert test == 89999999

0 comments on commit 8db27a3

Please sign in to comment.