adding tests for single token mint case and rewriting jit classifier to handle uni-v3 single token liquidity

mev_inspect/jit_liquidity.py

@@ -10,6 +10,7 @@ from mev_inspect.schemas.traces import (
     DecodedCallTrace,
     Protocol,
 )
+from mev_inspect.schemas.transfers import Transfer
 from mev_inspect.traces import is_child_trace_address
 from mev_inspect.transfers import get_net_transfers
 
@@ -17,6 +18,8 @@ LIQUIDITY_MINT_ROUTERS = [
     "0xC36442b4a4522E871399CD717aBDD847Ab11FE88".lower(),  # Uniswap V3 NFT Position Manager
 ]
 
+ZERO_ADDRESS = "0x0000000000000000000000000000000000000000"
+
 
 class JITTransferInfo(BaseModel):
     token0_address: str
@@ -179,22 +182,30 @@ def _get_transfer_info(
     if len(mint_net_transfers) > 2 or len(burn_net_transfers) > 2:
         error_found = True
 
-    token0_address, token1_address = _get_token_order(
+    if (
-        mint_net_transfers[0].token_address, mint_net_transfers[1].token_address
+        len(mint_net_transfers) < 2 or len(burn_net_transfers) < 2
-    )
+    ):  # Uniswap V3 Limit Case
-    if mint_net_transfers[0].token_address == token0_address:
+        if len(mint_net_transfers) == 0 or len(burn_net_transfers) == 0:
-        mint_token0 = mint_net_transfers[0].amount
+            raise Exception(
-        mint_token1 = mint_net_transfers[1].amount
+                "JIT Liquidity found where no tokens are transferred to pool address"
-    else:
+            )
-        mint_token0 = mint_net_transfers[1].amount
+
-        mint_token1 = mint_net_transfers[0].amount
+        return _parse_liquidity_limit_order(
+            mint_net_transfers, burn_net_transfers, error_found
+        )
 
-    if burn_net_transfers[0].token_address == token0_address:
-        burn_token0 = burn_net_transfers[0].amount
-        burn_token1 = burn_net_transfers[1].amount
     else:
-        burn_token0 = burn_net_transfers[1].amount
+        token0_address, token1_address = _get_token_order(
-        burn_token1 = burn_net_transfers[0].amount
+            mint_net_transfers[0].token_address, mint_net_transfers[1].token_address
+        )
+
+        mint_token0, mint_token1 = _parse_token_amounts(
+            token0_address, mint_net_transfers
+        )
+
+        burn_token0, burn_token1 = _parse_token_amounts(
+            token0_address, burn_net_transfers
+        )
 
     return JITTransferInfo(
         token0_address=token0_address,
@@ -226,11 +237,72 @@ def _get_bot_address(
             ):
                 return _get_bot_address(bot_trace[0], classified_traces)
             else:
-                return "0x0000000000000000000000000000000000000000"
+                return ZERO_ADDRESS
 
         elif type(mint_trace.from_address) == str:
             return mint_trace.from_address
+
+    return ZERO_ADDRESS
+
+
+def _parse_liquidity_limit_order(
+    mint_net_transfers: List[Transfer],
+    burn_net_transfers: List[Transfer],
+    error_found: bool,
+) -> JITTransferInfo:
+    try:
+        token0_address, token1_address = _get_token_order(
+            burn_net_transfers[0].token_address, burn_net_transfers[1].token_address
+        )
+    except IndexError:
+        token0_address, token1_address = _get_token_order(
+            mint_net_transfers[0].token_address, mint_net_transfers[1].token_address
+        )
+
+    if len(mint_net_transfers) < 2:
+        if token0_address == mint_net_transfers[0].token_address:
+            mint_token0 = mint_net_transfers[0].amount
+            mint_token1 = 0
         else:
-            return "0x0000000000000000000000000000000000000000"
+            mint_token0 = 0
+            mint_token1 = mint_net_transfers[0].amount
+
+        burn_token0, burn_token1 = _parse_token_amounts(
+            token0_address, burn_net_transfers
+        )
+
     else:
-        return "0x0000000000000000000000000000000000000000"
+        if token0_address == burn_net_transfers[0].token_address:
+            burn_token0 = burn_net_transfers[0].amount
+            burn_token1 = 0
+        else:
+            burn_token0 = 0
+            burn_token1 = burn_net_transfers[0].amount
+
+        mint_token0, mint_token1 = _parse_token_amounts(
+            token0_address, mint_net_transfers
+        )
+
+    return JITTransferInfo(
+        token0_address=token0_address,
+        token1_address=token1_address,
+        mint_token0=mint_token0,
+        mint_token1=mint_token1,
+        burn_token0=burn_token0,
+        burn_token1=burn_token1,
+        error=error_found,
+    )
+
+
+def _parse_token_amounts(
+    token0_address: str, net_transfers: List[Transfer]
+) -> Tuple[int, int]:
+    if token0_address == net_transfers[0].token_address:
+        token0_amount = net_transfers[0].amount
+        token1_amount = net_transfers[1].amount
+
+    else:
+        token0_amount = net_transfers[1].amount
+        token1_amount = net_transfers[0].amount
+
+    return token0_amount, token1_amount

tests/test_jit_liquidity.py

@@ -96,3 +96,48 @@ def test_single_sandwich_jit_liquidity_CRV_WETH(trace_classifier: TraceClassifie
     ]
 
     assert jit_liquidity_instances == expected_jit_liquidity
+
+
+def test_single_mint_token_jit(trace_classifier):
+    test_block = load_test_block(14643923)
+    classified_traces = trace_classifier.classify(test_block.traces)
+    swaps = get_swaps(classified_traces)
+    jit_liquidity_instances = get_jit_liquidity(classified_traces, swaps)
+
+    jit_swap = Swap(  # Double check these values
+        abi_name="UniswapV3Pool",
+        transaction_hash="0x43f9656e051a8e3b37f66668851922c6e8e4749d5a7aad605f21119cde541e49".lower(),
+        transaction_position=4,
+        block_number=14643923,
+        trace_address=[1, 0, 1, 0, 1, 0, 0],
+        contract_address="0xac4b3dacb91461209ae9d41ec517c2b9cb1b7daf".lower(),
+        from_address="0x74de5d4fcbf63e00296fd95d33236b9794016631".lower(),
+        to_address="0xdef1c0ded9bec7f1a1670819833240f027b25eff".lower(),
+        token_in_address="0x4d224452801aced8b2f0aebe155379bb5d594381".lower(),  # USDC Contract
+        token_in_amount=6522531010660457256888,
+        token_out_address="0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2".lower(),
+        token_out_amount=36485453136086109896,
+        protocol=Protocol.uniswap_v3,
+    )
+    expected_jit_liquidity = [
+        JITLiquidity(
+            block_number=14643923,
+            bot_address="0xa57Bd00134B2850B2a1c55860c9e9ea100fDd6CF".lower(),
+            pool_address="0xac4b3dacb91461209ae9d41ec517c2b9cb1b7daf".lower(),
+            mint_transaction_hash="0x003e36cb5d78924c5beaeef15db00cad94009856fe483a031d52ae975557ef53".lower(),
+            mint_trace=[0, 7, 1],
+            burn_transaction_hash="0xec9b2988f6c88968250c3904f6d2d6573f7284cb422b8022a14b7f0dac546348".lower(),
+            burn_trace=[0, 1, 0],
+            swaps=[jit_swap],
+            token0_address="0x4d224452801aced8b2f0aebe155379bb5d594381".lower(),
+            token1_address="0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2".lower(),
+            mint_token0_amount=0,
+            mint_token1_amount=9073930631365320229693,
+            burn_token0_amount=2424427669988518000798,
+            burn_token1_amount=9060377725722224517671,
+            token0_swap_volume=6522531010660457256888,
+            token1_swap_volume=0,
+        )
+    ]
+
+    assert jit_liquidity_instances == expected_jit_liquidity