Update module name to github.com/fraktal/mev-beta and fix channel closing issues in pipeline stages

This commit is contained in:
Krypto Kajun
2025-09-12 19:08:38 -05:00
parent fbb85e529a
commit 1113d82499
31 changed files with 3359 additions and 210 deletions

View File

@@ -0,0 +1,129 @@
package config
import (
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLoad(t *testing.T) {
// Create a temporary config file for testing
tmpFile, err := os.CreateTemp("", "config_test_*.yaml")
require.NoError(t, err)
defer os.Remove(tmpFile.Name())
// Write test config content
configContent := `
arbitrum:
rpc_endpoint: "https://arb1.arbitrum.io/rpc"
ws_endpoint: ""
chain_id: 42161
rate_limit:
requests_per_second: 10
max_concurrent: 5
burst: 20
bot:
enabled: true
polling_interval: 1
min_profit_threshold: 10.0
gas_price_multiplier: 1.2
max_workers: 10
channel_buffer_size: 100
rpc_timeout: 30
uniswap:
factory_address: "0x1F98431c8aD98523631AE4a59f267346ea31F984"
position_manager_address: "0xC36442b4a4522E871399CD717aBDD847Ab11FE88"
fee_tiers:
- 500
- 3000
- 10000
cache:
enabled: true
expiration: 300
max_size: 10000
log:
level: "info"
format: "text"
file: ""
database:
file: "mev-bot.db"
max_open_connections: 10
max_idle_connections: 5
`
_, err = tmpFile.Write([]byte(configContent))
require.NoError(t, err)
err = tmpFile.Close()
require.NoError(t, err)
// Test loading the config
cfg, err := Load(tmpFile.Name())
require.NoError(t, err)
// Verify the loaded config
assert.Equal(t, "https://arb1.arbitrum.io/rpc", cfg.Arbitrum.RPCEndpoint)
assert.Equal(t, int64(42161), cfg.Arbitrum.ChainID)
assert.Equal(t, 10, cfg.Arbitrum.RateLimit.RequestsPerSecond)
assert.True(t, cfg.Bot.Enabled)
assert.Equal(t, 1, cfg.Bot.PollingInterval)
assert.Equal(t, 10.0, cfg.Bot.MinProfitThreshold)
assert.Equal(t, "0x1F98431c8aD98523631AE4a59f267346ea31F984", cfg.Uniswap.FactoryAddress)
assert.Len(t, cfg.Uniswap.FeeTiers, 3)
assert.Equal(t, true, cfg.Uniswap.Cache.Enabled)
assert.Equal(t, "info", cfg.Log.Level)
assert.Equal(t, "mev-bot.db", cfg.Database.File)
}
func TestLoadWithInvalidFile(t *testing.T) {
// Test loading a non-existent config file
_, err := Load("/non/existent/file.yaml")
assert.Error(t, err)
}
func TestOverrideWithEnv(t *testing.T) {
// Create a temporary config file for testing
tmpFile, err := os.CreateTemp("", "config_test_*.yaml")
require.NoError(t, err)
defer os.Remove(tmpFile.Name())
// Write test config content
configContent := `
arbitrum:
rpc_endpoint: "https://arb1.arbitrum.io/rpc"
rate_limit:
requests_per_second: 10
max_concurrent: 5
bot:
max_workers: 10
channel_buffer_size: 100
`
_, err = tmpFile.Write([]byte(configContent))
require.NoError(t, err)
err = tmpFile.Close()
require.NoError(t, err)
// Set environment variables to override config
os.Setenv("ARBITRUM_RPC_ENDPOINT", "https://override.arbitrum.io/rpc")
os.Setenv("RPC_REQUESTS_PER_SECOND", "20")
os.Setenv("BOT_MAX_WORKERS", "20")
defer func() {
os.Unsetenv("ARBITRUM_RPC_ENDPOINT")
os.Unsetenv("RPC_REQUESTS_PER_SECOND")
os.Unsetenv("BOT_MAX_WORKERS")
}()
// Load the config
cfg, err := Load(tmpFile.Name())
require.NoError(t, err)
// Verify the overridden values
assert.Equal(t, "https://override.arbitrum.io/rpc", cfg.Arbitrum.RPCEndpoint)
assert.Equal(t, 20, cfg.Arbitrum.RateLimit.RequestsPerSecond)
assert.Equal(t, 20, cfg.Bot.MaxWorkers)
}

View File

@@ -0,0 +1,245 @@
package logger
import (
"bytes"
"io"
"os"
"testing"
"github.com/stretchr/testify/assert"
)
func TestNewLogger(t *testing.T) {
// Test creating a logger with stdout
logger := New("info", "text", "")
assert.NotNil(t, logger)
assert.NotNil(t, logger.logger)
assert.Equal(t, "info", logger.level)
}
func TestNewLoggerWithFile(t *testing.T) {
// Create a temporary file for testing
tmpFile, err := os.CreateTemp("", "logger_test_*.log")
assert.NoError(t, err)
defer os.Remove(tmpFile.Name())
err = tmpFile.Close()
assert.NoError(t, err)
// Test creating a logger with a file
logger := New("info", "text", tmpFile.Name())
assert.NotNil(t, logger)
assert.NotNil(t, logger.logger)
assert.Equal(t, "info", logger.level)
}
func TestDebug(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with debug level
logger := New("debug", "text", "")
// Log a debug message
logger.Debug("test debug message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the debug message
assert.Contains(t, output, "DEBUG:")
assert.Contains(t, output, "test debug message")
}
func TestDebugWithInfoLevel(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with info level (should not log debug messages)
logger := New("info", "text", "")
// Log a debug message
logger.Debug("test debug message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output does not contain the debug message
assert.NotContains(t, output, "DEBUG:")
assert.NotContains(t, output, "test debug message")
}
func TestInfo(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with info level
logger := New("info", "text", "")
// Log an info message
logger.Info("test info message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the info message
assert.Contains(t, output, "INFO:")
assert.Contains(t, output, "test info message")
}
func TestInfoWithDebugLevel(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with debug level
logger := New("debug", "text", "")
// Log an info message
logger.Info("test info message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the info message
assert.Contains(t, output, "INFO:")
assert.Contains(t, output, "test info message")
}
func TestWarn(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with warn level
logger := New("warn", "text", "")
// Log a warning message
logger.Warn("test warn message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the warning message
assert.Contains(t, output, "WARN:")
assert.Contains(t, output, "test warn message")
}
func TestWarnWithInfoLevel(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with info level (should log warnings)
logger := New("info", "text", "")
// Log a warning message
logger.Warn("test warn message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the warning message
assert.Contains(t, output, "WARN:")
assert.Contains(t, output, "test warn message")
}
func TestError(t *testing.T) {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger
logger := New("error", "text", "")
// Log an error message
logger.Error("test error message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the error message
assert.Contains(t, output, "ERROR:")
assert.Contains(t, output, "test error message")
}
func TestErrorWithAllLevels(t *testing.T) {
// Test that error messages are logged at all levels
levels := []string{"debug", "info", "warn", "error"}
for _, level := range levels {
// Capture stdout
old := os.Stdout
r, w, _ := os.Pipe()
os.Stdout = w
// Create logger with current level
logger := New(level, "text", "")
// Log an error message
logger.Error("test error message")
// Restore stdout
w.Close()
os.Stdout = old
// Read the captured output
var buf bytes.Buffer
io.Copy(&buf, r)
output := buf.String()
// Verify the output contains the error message
assert.Contains(t, output, "ERROR:")
assert.Contains(t, output, "test error message")
}
}

View File

@@ -5,7 +5,7 @@ import (
"fmt"
"sync"
"github.com/your-username/mev-beta/internal/config"
"github.com/fraktal/mev-beta/internal/config"
"golang.org/x/time/rate"
)

View File

@@ -0,0 +1,233 @@
package ratelimit
import (
"context"
"testing"
"time"
"github.com/fraktal/mev-beta/internal/config"
"github.com/stretchr/testify/assert"
"golang.org/x/time/rate"
)
func TestNewLimiterManager(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
FallbackEndpoints: []config.EndpointConfig{
{
URL: "https://fallback.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 5,
Burst: 10,
},
},
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Verify limiter manager was created correctly
assert.NotNil(t, lm)
assert.NotNil(t, lm.limiters)
assert.Len(t, lm.limiters, 2) // Primary + 1 fallback
// Check primary endpoint limiter
primaryLimiter, exists := lm.limiters[cfg.RPCEndpoint]
assert.True(t, exists)
assert.Equal(t, cfg.RPCEndpoint, primaryLimiter.URL)
assert.Equal(t, cfg.RateLimit, primaryLimiter.Config)
assert.NotNil(t, primaryLimiter.Limiter)
// Check fallback endpoint limiter
fallbackLimiter, exists := lm.limiters[cfg.FallbackEndpoints[0].URL]
assert.True(t, exists)
assert.Equal(t, cfg.FallbackEndpoints[0].URL, fallbackLimiter.URL)
assert.Equal(t, cfg.FallbackEndpoints[0].RateLimit, fallbackLimiter.Config)
assert.NotNil(t, fallbackLimiter.Limiter)
}
func TestWaitForLimit(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Test waiting for limit on existing endpoint
ctx := context.Background()
err := lm.WaitForLimit(ctx, cfg.RPCEndpoint)
assert.NoError(t, err)
// Test waiting for limit on non-existing endpoint
err = lm.WaitForLimit(ctx, "https://nonexistent.com")
assert.Error(t, err)
assert.Contains(t, err.Error(), "no rate limiter found for endpoint")
}
func TestTryWaitForLimit(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Test trying to wait for limit on existing endpoint
ctx := context.Background()
err := lm.TryWaitForLimit(ctx, cfg.RPCEndpoint)
assert.NoError(t, err) // Should succeed since we have burst capacity
// Test trying to wait for limit on non-existing endpoint
err = lm.TryWaitForLimit(ctx, "https://nonexistent.com")
assert.Error(t, err)
assert.Contains(t, err.Error(), "no rate limiter found for endpoint")
}
func TestGetLimiter(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Test getting limiter for existing endpoint
limiter, err := lm.GetLimiter(cfg.RPCEndpoint)
assert.NoError(t, err)
assert.NotNil(t, limiter)
assert.IsType(t, &rate.Limiter{}, limiter)
// Test getting limiter for non-existing endpoint
limiter, err = lm.GetLimiter("https://nonexistent.com")
assert.Error(t, err)
assert.Contains(t, err.Error(), "no rate limiter found for endpoint")
assert.Nil(t, limiter)
}
func TestUpdateLimiter(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Get original limiter
originalLimiter, err := lm.GetLimiter(cfg.RPCEndpoint)
assert.NoError(t, err)
assert.NotNil(t, originalLimiter)
// Update the limiter
newConfig := config.RateLimitConfig{
RequestsPerSecond: 20,
Burst: 40,
}
lm.UpdateLimiter(cfg.RPCEndpoint, newConfig)
// Get updated limiter
updatedLimiter, err := lm.GetLimiter(cfg.RPCEndpoint)
assert.NoError(t, err)
assert.NotNil(t, updatedLimiter)
// The limiter should be different (new instance)
assert.NotEqual(t, originalLimiter, updatedLimiter)
// Check that the config was updated
endpointLimiter := lm.limiters[cfg.RPCEndpoint]
assert.Equal(t, newConfig, endpointLimiter.Config)
}
func TestGetEndpoints(t *testing.T) {
// Create test config
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 10,
Burst: 20,
},
FallbackEndpoints: []config.EndpointConfig{
{
URL: "https://fallback1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 5,
Burst: 10,
},
},
{
URL: "https://fallback2.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 3,
Burst: 6,
},
},
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Get endpoints
endpoints := lm.GetEndpoints()
// Verify results
assert.Len(t, endpoints, 3) // Primary + 2 fallbacks
assert.Contains(t, endpoints, cfg.RPCEndpoint)
assert.Contains(t, endpoints, cfg.FallbackEndpoints[0].URL)
assert.Contains(t, endpoints, cfg.FallbackEndpoints[1].URL)
}
func TestRateLimiting(t *testing.T) {
// Create test config with very low rate limit for testing
cfg := &config.ArbitrumConfig{
RPCEndpoint: "https://arb1.arbitrum.io/rpc",
RateLimit: config.RateLimitConfig{
RequestsPerSecond: 1, // 1 request per second
Burst: 1, // No burst
},
}
// Create limiter manager
lm := NewLimiterManager(cfg)
// Make a request (should succeed immediately)
start := time.Now()
ctx := context.Background()
err := lm.WaitForLimit(ctx, cfg.RPCEndpoint)
assert.NoError(t, err)
duration := time.Since(start)
assert.True(t, duration < time.Millisecond*100, "First request should be fast")
// Make another request immediately (should be delayed)
start = time.Now()
err = lm.WaitForLimit(ctx, cfg.RPCEndpoint)
assert.NoError(t, err)
duration = time.Since(start)
assert.True(t, duration >= time.Second, "Second request should be delayed by rate limiter")
}