mev-beta/internal/ratelimit/manager_test.go

package ratelimit

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"golang.org/x/time/rate"

	"github.com/fraktal/mev-beta/internal/config"
)

func TestNewLimiterManager(t *testing.T) {
	// Create test config
	cfg := &config.ArbitrumConfig{
		RPCEndpoint: "https://arb1.arbitrum.io/rpc",
		RateLimit: config.RateLimitConfig{
			RequestsPerSecond: 10,
			Burst:             20,
		},
		ReadingEndpoints: []config.EndpointConfig{
			{
				URL: "https://read.arbitrum.io/rpc",
				RateLimit: config.RateLimitConfig{
					RequestsPerSecond: 5,
					Burst:             10,
				},
			},
		},
		ExecutionEndpoints: []config.EndpointConfig{
			{
				URL: "https://exec.arbitrum.io/rpc",
				RateLimit: config.RateLimitConfig{
					RequestsPerSecond: 3,
					Burst:             6,
				},
			},
		},
	}

	// 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, 3) // 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.ReadingEndpoints[0].URL]
	assert.True(t, exists)
	assert.Equal(t, cfg.ReadingEndpoints[0].URL, fallbackLimiter.URL)
	assert.Equal(t, cfg.ReadingEndpoints[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,
		},
		ReadingEndpoints: []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.ReadingEndpoints[0].URL)
	assert.Contains(t, endpoints, cfg.ReadingEndpoints[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")
}