Files
mev-beta/pkg/arbitrage/multihop_test.go
2025-09-16 11:05:47 -05:00

404 lines
12 KiB
Go

package arbitrage
import (
"context"
"math/big"
"testing"
"time"
"github.com/ethereum/go-ethereum/common"
"github.com/fraktal/mev-beta/internal/logger"
"github.com/fraktal/mev-beta/pkg/market"
"github.com/holiman/uint256"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
// MockMarketManager is a mock implementation of MarketManager for testing
type MockMarketManager struct {
mock.Mock
}
func (m *MockMarketManager) GetAllPools() []market.PoolData {
args := m.Called()
return args.Get(0).([]market.PoolData)
}
func (m *MockMarketManager) GetPool(ctx context.Context, poolAddress common.Address) (*market.PoolData, error) {
args := m.Called(ctx, poolAddress)
if args.Get(0) == nil {
return nil, args.Error(1)
}
return args.Get(0).(*market.PoolData), args.Error(1)
}
func (m *MockMarketManager) GetPoolsByTokens(token0, token1 common.Address) []*market.PoolData {
args := m.Called(token0, token1)
return args.Get(0).([]*market.PoolData)
}
func (m *MockMarketManager) UpdatePool(poolAddress common.Address, liquidity *uint256.Int, sqrtPriceX96 *uint256.Int, tick int) {
m.Called(poolAddress, liquidity, sqrtPriceX96, tick)
}
func (m *MockMarketManager) GetPoolsByTokensWithProtocol(token0, token1 common.Address, protocol string) []*market.PoolData {
args := m.Called(token0, token1, protocol)
return args.Get(0).([]*market.PoolData)
}
// TestNewMultiHopScanner tests the creation of a new MultiHopScanner
func TestNewMultiHopScanner(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
assert.NotNil(t, scanner)
assert.Equal(t, log, scanner.logger)
assert.Equal(t, marketMgr, scanner.marketMgr)
assert.Equal(t, 4, scanner.maxHops)
assert.Equal(t, "1000000000000000", scanner.minProfitWei.String())
assert.Equal(t, 0.03, scanner.maxSlippage)
assert.Equal(t, 100, scanner.maxPaths)
assert.Equal(t, time.Millisecond*500, scanner.pathTimeout)
assert.NotNil(t, scanner.pathCache)
assert.NotNil(t, scanner.tokenGraph)
assert.NotNil(t, scanner.pools)
}
// TestTokenGraph tests the TokenGraph functionality
func TestTokenGraph(t *testing.T) {
graph := NewTokenGraph()
assert.NotNil(t, graph)
assert.NotNil(t, graph.adjacencyList)
// Test adding edges
tokenA := common.HexToAddress("0xA")
tokenB := common.HexToAddress("0xB")
pool := &PoolInfo{
Address: common.HexToAddress("0x1"),
Token0: tokenA,
Token1: tokenB,
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
}
// Add pool to graph
graph.mutex.Lock()
graph.adjacencyList[tokenA] = make(map[common.Address][]*PoolInfo)
graph.adjacencyList[tokenA][tokenB] = append(graph.adjacencyList[tokenA][tokenB], pool)
graph.mutex.Unlock()
// Test getting adjacent tokens
adjacent := graph.GetAdjacentTokens(tokenA)
assert.Len(t, adjacent, 1)
assert.Contains(t, adjacent, tokenB)
assert.Len(t, adjacent[tokenB], 1)
assert.Equal(t, pool, adjacent[tokenB][0])
}
// TestIsPoolUsable tests the isPoolUsable function
func TestIsPoolUsable(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Test usable pool (recent and sufficient liquidity)
now := time.Now()
usablePool := &PoolInfo{
Address: common.HexToAddress("0x1"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000), // 1 ETH worth of liquidity
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: now,
}
assert.True(t, scanner.isPoolUsable(usablePool))
// Test pool with insufficient liquidity
unusablePool1 := &PoolInfo{
Address: common.HexToAddress("0x2"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(10000000000000000), // 0.01 ETH worth of liquidity (too little)
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: now,
}
assert.False(t, scanner.isPoolUsable(unusablePool1))
// Test stale pool
stalePool := &PoolInfo{
Address: common.HexToAddress("0x3"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: now.Add(-10 * time.Minute), // 10 minutes ago (stale)
}
assert.False(t, scanner.isPoolUsable(stalePool))
}
// TestCalculateSimpleAMMOutput tests the calculateSimpleAMMOutput function
func TestCalculateSimpleAMMOutput(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Create a pool with known values for testing
tokenIn := common.HexToAddress("0xA")
tokenOut := common.HexToAddress("0xB")
// Create a pool with realistic values
// SqrtPriceX96 = 79228162514264337593543950336 (represents 1.0 price)
// Liquidity = 1000000000000000000 (1 ETH)
pool := &PoolInfo{
Address: common.HexToAddress("0x1"),
Token0: tokenIn,
Token1: tokenOut,
Protocol: "UniswapV2",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
}
// Test with 1 ETH input
amountIn := big.NewInt(1000000000000000000) // 1 ETH
output, err := scanner.calculateSimpleAMMOutput(amountIn, pool, tokenIn, tokenOut)
// We should get a valid output
assert.NoError(t, err)
assert.NotNil(t, output)
assert.True(t, output.Sign() > 0)
// Test with missing data
badPool := &PoolInfo{
Address: common.HexToAddress("0x2"),
Token0: tokenIn,
Token1: tokenOut,
Protocol: "UniswapV2",
Fee: 3000,
Liquidity: nil, // Missing liquidity
SqrtPriceX96: nil, // Missing sqrtPriceX96
LastUpdated: time.Now(),
}
output, err = scanner.calculateSimpleAMMOutput(amountIn, badPool, tokenIn, tokenOut)
assert.Error(t, err)
assert.Nil(t, output)
}
// TestCalculateUniswapV3Output tests the calculateUniswapV3Output function
func TestCalculateUniswapV3Output(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Create a pool with known values for testing
tokenIn := common.HexToAddress("0xA")
tokenOut := common.HexToAddress("0xB")
// Create a pool with realistic values
pool := &PoolInfo{
Address: common.HexToAddress("0x1"),
Token0: tokenIn,
Token1: tokenOut,
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
}
// Test with 1 ETH input
amountIn := big.NewInt(1000000000000000000) // 1 ETH
output, err := scanner.calculateUniswapV3Output(amountIn, pool, tokenIn, tokenOut)
// We should get a valid output
assert.NoError(t, err)
assert.NotNil(t, output)
assert.True(t, output.Sign() > 0)
}
// TestEstimateHopGasCost tests the estimateHopGasCost function
func TestEstimateHopGasCost(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Test UniswapV3
gas := scanner.estimateHopGasCost("UniswapV3")
assert.Equal(t, int64(150000), gas.Int64())
// Test UniswapV2
gas = scanner.estimateHopGasCost("UniswapV2")
assert.Equal(t, int64(120000), gas.Int64())
// Test SushiSwap
gas = scanner.estimateHopGasCost("SushiSwap")
assert.Equal(t, int64(120000), gas.Int64())
// Test default case
gas = scanner.estimateHopGasCost("UnknownProtocol")
assert.Equal(t, int64(150000), gas.Int64())
}
// TestIsProfitable tests the isProfitable function
func TestIsProfitable(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Create a profitable path
profitablePath := &ArbitragePath{
NetProfit: big.NewInt(2000000000000000000), // 2 ETH profit
ROI: 5.0, // 5% ROI
}
assert.True(t, scanner.isProfitable(profitablePath))
// Create an unprofitable path (below minimum profit)
unprofitablePath1 := &ArbitragePath{
NetProfit: big.NewInt(100000000000000000), // 0.1 ETH profit (below 0.001 ETH threshold)
ROI: 0.5, // 0.5% ROI
}
assert.False(t, scanner.isProfitable(unprofitablePath1))
// Create a path with good profit but poor ROI
unprofitablePath2 := &ArbitragePath{
NetProfit: big.NewInt(5000000000000000000), // 5 ETH profit
ROI: 0.5, // 0.5% ROI (below 1% threshold)
}
assert.False(t, scanner.isProfitable(unprofitablePath2))
}
// TestCreateArbitragePath tests the createArbitragePath function
func TestCreateArbitragePath(t *testing.T) {
log := logger.New("info", "text", "")
marketMgr := &market.MarketManager{}
scanner := NewMultiHopScanner(log, marketMgr)
// Test with invalid inputs
tokens := []common.Address{
common.HexToAddress("0xA"),
common.HexToAddress("0xB"),
}
pools := []*PoolInfo{
{
Address: common.HexToAddress("0x1"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
},
}
initialAmount := big.NewInt(1000000000000000000) // 1 ETH
// This should fail because we need at least 3 tokens for a valid arbitrage path (A->B->A)
path := scanner.createArbitragePath(tokens, pools, initialAmount)
assert.Nil(t, path)
// Test with valid inputs (triangle: A->B->C->A)
validTokens := []common.Address{
common.HexToAddress("0xA"),
common.HexToAddress("0xB"),
common.HexToAddress("0xC"),
common.HexToAddress("0xA"), // Back to start
}
validPools := []*PoolInfo{
{
Address: common.HexToAddress("0x1"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
},
{
Address: common.HexToAddress("0x2"),
Token0: common.HexToAddress("0xB"),
Token1: common.HexToAddress("0xC"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
},
{
Address: common.HexToAddress("0x3"),
Token0: common.HexToAddress("0xC"),
Token1: common.HexToAddress("0xA"),
Protocol: "UniswapV3",
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
},
}
path = scanner.createArbitragePath(validTokens, validPools, initialAmount)
assert.NotNil(t, path)
assert.Len(t, path.Tokens, 4)
assert.Len(t, path.Pools, 3)
assert.Len(t, path.Protocols, 3)
assert.Len(t, path.Fees, 3)
assert.NotNil(t, path.EstimatedGas)
assert.NotNil(t, path.NetProfit)
}
// TestScanForArbitrage tests the main ScanForArbitrage function
func TestScanForArbitrage(t *testing.T) {
log := logger.New("info", "text", "")
// Create a mock market manager
mockMarketMgr := &MockMarketManager{}
// Set up mock expectations
mockMarketMgr.On("GetAllPools").Return([]market.PoolData{
{
Address: common.HexToAddress("0x1"),
Token0: common.HexToAddress("0xA"),
Token1: common.HexToAddress("0xB"),
Fee: 3000,
Liquidity: uint256.NewInt(1000000000000000000),
SqrtPriceX96: uint256.NewInt(79228162514264337593543950336),
LastUpdated: time.Now(),
},
})
scanner := NewMultiHopScanner(log, mockMarketMgr)
ctx := context.Background()
triggerToken := common.HexToAddress("0xA")
amount := big.NewInt(1000000000000000000) // 1 ETH
paths, err := scanner.ScanForArbitrage(ctx, triggerToken, amount)
// For now, we expect it to return without error, even if no profitable paths are found
assert.NoError(t, err)
assert.NotNil(t, paths)
// It's okay to return an empty slice if no profitable paths are found
}