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

2025-09-12 19:08:38 -05:00
parent fbb85e529a
commit 1113d82499
31 changed files with 3359 additions and 210 deletions
--- a/internal/ratelimit/manager.go
+++ b/internal/ratelimit/manager.go
@@ -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"
 )

--- a/internal/ratelimit/manager_test.go
+++ b/internal/ratelimit/manager_test.go
@@ -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")
+}