mev-beta/pkg/market/pipeline_test.go

package market

import (
	"context"
	"math/big"
	"testing"

	"github.com/ethereum/go-ethereum/common"
	"github.com/fraktal/mev-beta/internal/config"
	"github.com/fraktal/mev-beta/internal/logger"
	"github.com/fraktal/mev-beta/pkg/events"
	scannerpkg "github.com/fraktal/mev-beta/pkg/scanner"
	"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) GetPool(ctx context.Context, poolAddress common.Address) (*PoolData, error) {
	args := m.Called(ctx, poolAddress)
	return args.Get(0).(*PoolData), args.Error(1)
}

func (m *MockMarketManager) GetPoolsByTokens(token0, token1 common.Address) []*PoolData {
	args := m.Called(token0, token1)
	return args.Get(0).([]*PoolData)
}

// MockLogger is a mock implementation of logger.Logger for testing
type MockLogger struct {
	mock.Mock
}

func (m *MockLogger) Debug(msg string) {
	m.Called(msg)
}

func (m *MockLogger) Info(msg string) {
	m.Called(msg)
}

func (m *MockLogger) Warn(msg string) {
	m.Called(msg)
}

func (m *MockLogger) Error(msg string, err ...interface{}) {
	m.Called(msg, err)
}

func TestNewPipeline(t *testing.T) {
	// Create mock config
	cfg := &config.BotConfig{
		MaxWorkers:        5,
		ChannelBufferSize: 10,
	}

	// Create mock logger
	logger := logger.New("info", "text", "")

	// Create mock market manager
	marketMgr := &MarketManager{}

	// Create mock scanner
	scannerObj := &scannerpkg.MarketScanner{}

	// Create pipeline
	pipeline := NewPipeline(cfg, logger, marketMgr, scannerObj, nil)

	// Verify pipeline was created correctly
	assert.NotNil(t, pipeline)
	assert.Equal(t, cfg, pipeline.config)
	assert.Equal(t, logger, pipeline.logger)
	assert.Equal(t, marketMgr, pipeline.marketMgr)
	assert.Equal(t, scannerObj, pipeline.scanner)
	assert.Equal(t, cfg.ChannelBufferSize, pipeline.bufferSize)
	assert.Equal(t, cfg.MaxWorkers, pipeline.concurrency)
	assert.NotNil(t, pipeline.eventParser)
	assert.Len(t, pipeline.stages, 1) // Should have TransactionDecoderStage by default
}

func TestAddStage(t *testing.T) {
	// Create pipeline
	cfg := &config.BotConfig{
		MaxWorkers:        5,
		ChannelBufferSize: 10,
	}
	logger := logger.New("info", "text", "")
	marketMgr := &MarketManager{}
	scannerObj := &scannerpkg.MarketScanner{}
	pipeline := NewPipeline(cfg, logger, marketMgr, scannerObj, nil)

	// Add a new stage
	newStage := func(ctx context.Context, input <-chan *events.Event, output chan<- *events.Event) error {
		return nil
	}
	pipeline.AddStage(newStage)

	// Verify stage was added
	assert.Len(t, pipeline.stages, 2) // TransactionDecoderStage + newStage
}

func TestAddDefaultStages(t *testing.T) {
	// Create pipeline
	cfg := &config.BotConfig{
		MaxWorkers:        5,
		ChannelBufferSize: 10,
	}
	logger := logger.New("info", "text", "")
	marketMgr := &MarketManager{}
	scannerObj := &scannerpkg.MarketScanner{}
	pipeline := NewPipeline(cfg, logger, marketMgr, scannerObj, nil)

	// Add default stages
	pipeline.AddDefaultStages()

	// Verify stages were added (should be 4 total: TransactionDecoder, MarketAnalysis, ArbitrageDetection, plus the initial TransactionDecoder)
	assert.Len(t, pipeline.stages, 4)
}

func TestTransactionDecoderStage(t *testing.T) {
	// Create mock config
	cfg := &config.BotConfig{
		MaxWorkers:        1, // Use 1 worker for simplicity in test
		ChannelBufferSize: 10,
	}

	// Create mock logger
	log := logger.New("info", "text", "")

	// Create mock market manager
	marketMgr := &MarketManager{}

	// Create the stage
	stage := TransactionDecoderStage(cfg, log, marketMgr, nil, nil)

	// Verify the stage function was created
	assert.NotNil(t, stage)
}

func TestCalculatePriceImpact(t *testing.T) {
	// Create test event
	event := &events.Event{
		Amount0: big.NewInt(1000000000), // 1000 tokens
		Amount1: big.NewInt(0),
	}

	// Create test pool data
	liquidity := uint256.NewInt(1000000000000000000) // 1 ETH in liquidity
	poolData := &PoolData{
		Liquidity: liquidity,
	}

	// Calculate price impact
	impact, err := calculatePriceImpact(event, poolData)

	// Verify results
	assert.NoError(t, err)
	assert.InDelta(t, 0.001, impact, 0.001) // 0.001% impact (1000/1000000000000000000 * 100)
}

func TestCalculatePriceImpactNoAmount(t *testing.T) {
	// Create test event with no amount
	event := &events.Event{
		Amount0: big.NewInt(0),
		Amount1: big.NewInt(0),
	}

	// Create test pool data
	liquidity := uint256.NewInt(10000000000000000000) // 10 ETH in liquidity
	poolData := &PoolData{
		Liquidity: liquidity,
	}

	// Calculate price impact
	impact, err := calculatePriceImpact(event, poolData)

	// Verify results
	assert.NoError(t, err)
	assert.Equal(t, 0.0, impact)
}

func TestCalculatePriceImpactNoLiquidity(t *testing.T) {
	// Create test event
	event := &events.Event{
		Amount0: big.NewInt(1000000000),
		Amount1: big.NewInt(0),
	}

	// Create test pool data with zero liquidity
	liquidity := uint256.NewInt(0)
	poolData := &PoolData{
		Liquidity: liquidity,
	}

	// Calculate price impact
	impact, err := calculatePriceImpact(event, poolData)

	// Verify results
	assert.NoError(t, err)
	assert.Equal(t, 0.0, impact)
}