mev-inspect-py/tests/test_arbitrages.py
2021-12-22 22:41:10 -05:00

254 lines
8.4 KiB
Python

from typing import List
from mev_inspect.arbitrages import _get_shortest_route, get_arbitrages
from mev_inspect.classifiers.specs.uniswap import (
UNISWAP_V2_PAIR_ABI_NAME,
UNISWAP_V3_POOL_ABI_NAME,
)
from mev_inspect.schemas.swaps import Swap
def test_two_pool_arbitrage(get_transaction_hashes, get_addresses):
block_number = 123
[transaction_hash] = get_transaction_hashes(1)
[
account_address,
first_pool_address,
second_pool_address,
unrelated_pool_address,
first_token_address,
second_token_address,
third_token_address,
] = get_addresses(7)
first_token_in_amount = 10
first_token_out_amount = 11
second_token_amount = 15
transaction_position = 0
arb_swaps = [
Swap(
abi_name=UNISWAP_V2_PAIR_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[0],
contract_address=first_pool_address,
from_address=account_address,
to_address=second_pool_address,
token_in_address=first_token_address,
token_in_amount=first_token_in_amount,
token_out_address=second_token_address,
token_out_amount=second_token_amount,
),
Swap(
abi_name=UNISWAP_V3_POOL_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[1],
contract_address=second_pool_address,
from_address=first_pool_address,
to_address=account_address,
token_in_address=second_token_address,
token_in_amount=second_token_amount,
token_out_address=first_token_address,
token_out_amount=first_token_out_amount,
),
]
unrelated_swap = Swap(
abi_name=UNISWAP_V3_POOL_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[2, 0],
contract_address=unrelated_pool_address,
from_address=account_address,
to_address=account_address,
token_in_address=second_token_address,
token_in_amount=first_token_in_amount,
token_out_address=third_token_address,
token_out_amount=first_token_out_amount,
)
swaps = [
unrelated_swap,
*arb_swaps,
]
arbitrages = get_arbitrages(swaps)
assert len(arbitrages) == 1
arbitrage = arbitrages[0]
assert arbitrage.swaps == arb_swaps
assert arbitrage.account_address == account_address
assert arbitrage.profit_token_address == first_token_address
assert arbitrage.start_amount == first_token_in_amount
assert arbitrage.end_amount == first_token_out_amount
assert arbitrage.profit_amount == first_token_out_amount - first_token_in_amount
def test_three_pool_arbitrage(get_transaction_hashes, get_addresses):
block_number = 123
[transaction_hash] = get_transaction_hashes(1)
[
account_address,
first_pool_address,
second_pool_address,
third_pool_address,
first_token_address,
second_token_address,
third_token_address,
] = get_addresses(7)
first_token_in_amount = 10
first_token_out_amount = 11
second_token_amount = 15
third_token_amount = 40
transaction_position = 0
swaps = [
Swap(
abi_name=UNISWAP_V2_PAIR_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[0],
contract_address=first_pool_address,
from_address=account_address,
to_address=second_pool_address,
token_in_address=first_token_address,
token_in_amount=first_token_in_amount,
token_out_address=second_token_address,
token_out_amount=second_token_amount,
),
Swap(
abi_name=UNISWAP_V3_POOL_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[1],
contract_address=second_pool_address,
from_address=first_pool_address,
to_address=third_pool_address,
token_in_address=second_token_address,
token_in_amount=second_token_amount,
token_out_address=third_token_address,
token_out_amount=third_token_amount,
),
Swap(
abi_name=UNISWAP_V3_POOL_ABI_NAME,
transaction_hash=transaction_hash,
transaction_position=transaction_position,
block_number=block_number,
trace_address=[2],
contract_address=third_pool_address,
from_address=second_pool_address,
to_address=account_address,
token_in_address=third_token_address,
token_in_amount=third_token_amount,
token_out_address=first_token_address,
token_out_amount=first_token_out_amount,
),
]
arbitrages = get_arbitrages(swaps)
assert len(arbitrages) == 1
arbitrage = arbitrages[0]
assert arbitrage.swaps == swaps
assert arbitrage.account_address == account_address
assert arbitrage.profit_token_address == first_token_address
assert arbitrage.start_amount == first_token_in_amount
assert arbitrage.end_amount == first_token_out_amount
assert arbitrage.profit_amount == first_token_out_amount - first_token_in_amount
def test_get_shortest_route():
# A -> B, B -> A
start_swap = create_generic_swap("0xa", "0xb")
end_swap = create_generic_swap("0xb", "0xa")
route = _get_shortest_route(start_swap, [end_swap], [])
assert route is not None
assert len(route) == 2
# A->B, B->C, C->A
start_swap = create_generic_swap("0xa", "0xb")
other_swaps = [create_generic_swap("0xb", "0xc")]
end_swap = create_generic_swap("0xc", "0xa")
route = _get_shortest_route(start_swap, [end_swap], other_swaps)
assert route is not None
assert len(route) == 3
# A->B, B->C, C->A + A->D
other_swaps.append(create_generic_swap("0xa", "0xd"))
route = _get_shortest_route(start_swap, [end_swap], other_swaps)
assert route is not None
assert len(route) == 3
# A->B, B->C, C->A + A->D B->E
other_swaps.append(create_generic_swap("0xb", "0xe"))
route = _get_shortest_route(start_swap, [end_swap], other_swaps)
assert route is not None
assert len(route) == 3
# A->B, B->A, B->C, C->A
other_swaps = [create_generic_swap("0xb", "0xa"), create_generic_swap("0xb", "0xc")]
route = _get_shortest_route(start_swap, [end_swap], other_swaps)
expected_smallest_route = [["0xa", "0xb"], ["0xb", "0xc"], ["0xc", "0xa"]]
assert route is not None
assert len(route) == len(expected_smallest_route)
for i, [expected_token_in, expected_token_out] in enumerate(
expected_smallest_route
):
assert expected_token_in == route[i].token_in_address
assert expected_token_out == route[i].token_out_address
# A->B, B->C, C->D, D->A, B->D
end_swap = create_generic_swap("0xd", "0xa")
other_swaps = [
create_generic_swap("0xb", "0xc"),
create_generic_swap("0xc", "0xd"),
create_generic_swap("0xb", "0xd"),
]
expected_smallest_route = [["0xa", "0xb"], ["0xb", "0xd"], ["0xd", "0xa"]]
route = _get_shortest_route(start_swap, [end_swap], other_swaps)
assert len(route) == 3
for i, [expected_token_in, expected_token_out] in enumerate(
expected_smallest_route
):
assert expected_token_in == route[i].token_in_address
assert expected_token_out == route[i].token_out_address
def create_generic_swap(
tok_a: str = "0xa",
tok_b: str = "0xb",
amount_a_in: int = 1,
amount_b_out: int = 1,
trace_address: List[int] = [],
):
return Swap(
abi_name=UNISWAP_V3_POOL_ABI_NAME,
transaction_hash="0xfake",
transaction_position=0,
block_number=0,
trace_address=trace_address,
contract_address="0xfake",
from_address="0xfake",
to_address="0xfake",
token_in_address=tok_a,
token_in_amount=amount_a_in,
token_out_address=tok_b,
token_out_amount=amount_b_out,
)