From eab6164c164bbd3185ba06e3b0ec7e4d2a0d96e4 Mon Sep 17 00:00:00 2001 From: Luke Van Seters Date: Mon, 9 Aug 2021 17:52:15 -0400 Subject: [PATCH] Add tests for swaps --- tests/helpers.py | 81 ++++++++++++++++++++++++++ tests/test_swaps.py | 135 +++++++++++++++++++++++++++++++++++++++++++ tests/test_traces.py | 37 ++---------- 3 files changed, 220 insertions(+), 33 deletions(-) create mode 100644 tests/helpers.py create mode 100644 tests/test_swaps.py diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..0ebcf0f --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,81 @@ +from typing import List + +from mev_inspect.schemas.blocks import TraceType +from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace + + +def make_transfer_trace( + block_number: int, + transaction_hash: str, + trace_address: List[int], + from_address: str, + to_address: str, + token_address: str, + amount: int, +): + return ClassifiedTrace( + transaction_hash=transaction_hash, + block_number=block_number, + trace_type=TraceType.call, + trace_address=trace_address, + classification=Classification.transfer, + from_address=from_address, + to_address=token_address, + inputs={ + "recipient": to_address, + "amount": amount, + }, + ) + + +def make_swap_trace( + block_number: int, + transaction_hash: str, + trace_address: List[int], + from_address: str, + pool_address: str, + abi_name: str, + recipient_address: str, + recipient_input_key: str, +): + return ClassifiedTrace( + transaction_hash=transaction_hash, + block_number=block_number, + trace_type=TraceType.call, + trace_address=trace_address, + classification=Classification.swap, + from_address=from_address, + to_address=pool_address, + inputs={recipient_input_key: recipient_address}, + abi_name=abi_name, + ) + + +def make_unknown_trace( + block_number, + transaction_hash, + trace_address, +): + return ClassifiedTrace( + transaction_hash=transaction_hash, + block_number=block_number, + trace_type=TraceType.call, + trace_address=trace_address, + classification=Classification.unknown, + ) + + +def make_many_unknown_traces( + block_number, + transaction_hash, + trace_addresses, +) -> List[ClassifiedTrace]: + + return [ + make_unknown_trace( + block_number, + transaction_hash, + trace_address, + ) + for trace_address in trace_addresses + ] diff --git a/tests/test_swaps.py b/tests/test_swaps.py new file mode 100644 index 0000000..bac65da --- /dev/null +++ b/tests/test_swaps.py @@ -0,0 +1,135 @@ +from mev_inspect.swaps import ( + get_swaps, + UNISWAP_V2_PAIR_ABI_NAME, + UNISWAP_V3_POOL_ABI_NAME, +) + +from .helpers import ( + make_unknown_trace, + make_transfer_trace, + make_swap_trace, +) + + +def test_swaps( + get_transaction_hashes, + get_addresses, +): + block_number = 123 + + [ + first_transaction_hash, + second_transaction_hash, + ] = get_transaction_hashes(2) + + [ + alice_address, + bob_address, + carl_address, + first_token_in_address, + first_token_out_address, + first_pool_address, + second_token_in_address, + second_token_out_address, + second_pool_address, + ] = get_addresses(9) + + first_token_in_amount = 10 + first_token_out_amount = 20 + second_token_in_amount = 30 + second_token_out_amount = 40 + + traces = [ + make_unknown_trace(block_number, first_transaction_hash, []), + make_transfer_trace( + block_number, + first_transaction_hash, + trace_address=[0], + from_address=alice_address, + to_address=first_pool_address, + token_address=first_token_in_address, + amount=first_token_in_amount, + ), + make_swap_trace( + block_number, + first_transaction_hash, + trace_address=[1], + from_address=alice_address, + pool_address=first_pool_address, + abi_name=UNISWAP_V2_PAIR_ABI_NAME, + recipient_address=bob_address, + recipient_input_key="to", + ), + make_transfer_trace( + block_number, + first_transaction_hash, + trace_address=[1, 0], + from_address=first_pool_address, + to_address=bob_address, + token_address=first_token_out_address, + amount=first_token_out_amount, + ), + make_swap_trace( + block_number, + second_transaction_hash, + trace_address=[], + from_address=bob_address, + pool_address=second_pool_address, + abi_name=UNISWAP_V3_POOL_ABI_NAME, + recipient_address=carl_address, + recipient_input_key="recipient", + ), + make_transfer_trace( + block_number, + second_transaction_hash, + trace_address=[0], + from_address=second_pool_address, + to_address=carl_address, + token_address=second_token_out_address, + amount=second_token_out_amount, + ), + make_transfer_trace( + block_number, + second_transaction_hash, + trace_address=[1], + from_address=bob_address, + to_address=second_pool_address, + token_address=second_token_in_address, + amount=second_token_in_amount, + ), + ] + + swaps = get_swaps(traces) + + assert len(swaps) == 2 + + if swaps[0].abi_name == UNISWAP_V2_PAIR_ABI_NAME: + [first_swap, second_swap] = swaps # pylint: disable=unbalanced-tuple-unpacking + else: + [second_swap, first_swap] = swaps # pylint: disable=unbalanced-tuple-unpacking + + assert first_swap.abi_name == UNISWAP_V2_PAIR_ABI_NAME + assert first_swap.transaction_hash == first_transaction_hash + assert first_swap.block_number == block_number + assert first_swap.trace_address == [1] + assert first_swap.protocol is None + assert first_swap.pool_address == first_pool_address + assert first_swap.from_address == alice_address + assert first_swap.to_address == bob_address + assert first_swap.token_in_address == first_token_in_address + assert first_swap.token_in_amount == first_token_in_amount + assert first_swap.token_out_address == first_token_out_address + assert first_swap.token_out_amount == first_token_out_amount + + assert second_swap.abi_name == UNISWAP_V3_POOL_ABI_NAME + assert second_swap.transaction_hash == second_transaction_hash + assert second_swap.block_number == block_number + assert second_swap.trace_address == [] + assert second_swap.protocol is None + assert second_swap.pool_address == second_pool_address + assert second_swap.from_address == bob_address + assert second_swap.to_address == carl_address + assert second_swap.token_in_address == second_token_in_address + assert second_swap.token_in_amount == second_token_in_amount + assert second_swap.token_out_address == second_token_out_address + assert second_swap.token_out_amount == second_token_out_amount diff --git a/tests/test_traces.py b/tests/test_traces.py index 489d677..4eca518 100644 --- a/tests/test_traces.py +++ b/tests/test_traces.py @@ -1,38 +1,9 @@ from typing import List -from mev_inspect.schemas.blocks import TraceType -from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace +from mev_inspect.schemas.classified_traces import ClassifiedTrace from mev_inspect.traces import is_child_trace_address, get_child_traces - -def make_unknown_classified_trace( - block_number, - transaction_hash, - trace_address, -): - return ClassifiedTrace( - transaction_hash=transaction_hash, - block_number=block_number, - trace_type=TraceType.call, - trace_address=trace_address, - classification=Classification.unknown, - ) - - -def make_traces( - block_number, - transaction_hash, - trace_addresses, -) -> List[ClassifiedTrace]: - - return [ - make_unknown_classified_trace( - block_number, - transaction_hash, - trace_address, - ) - for trace_address in trace_addresses - ] +from .helpers import make_many_unknown_traces def test_is_child_trace_address(): @@ -68,13 +39,13 @@ def test_get_child_traces(get_transaction_hashes): second_hash_trace_addresses = [[], [0], [1], [1, 0], [2]] - traces += make_traces( + traces += make_many_unknown_traces( block_number, first_hash, first_hash_trace_addresses, ) - traces += make_traces( + traces += make_many_unknown_traces( block_number, second_hash, second_hash_trace_addresses,