x/handler/router/observe_test.go

package router

import (
	"context"
	"errors"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/go-gost/core/handler"
	"github.com/go-gost/core/limiter/rate"
	"github.com/go-gost/core/observer"
	"github.com/go-gost/core/observer/stats"
	stats_util "github.com/go-gost/x/internal/util/stats"
)

// ---------------------------------------------------------------------------
// checkRateLimit tests
// ---------------------------------------------------------------------------

func TestCheckRateLimit_NilLimiter(t *testing.T) {
	h := &routerHandler{
		options: handler.Options{
			RateLimiter: nil,
		},
	}
	addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}
	if !h.checkRateLimit(addr) {
		t.Error("checkRateLimit returned false, want true")
	}
}

func TestCheckRateLimit_Allowed(t *testing.T) {
	h := &routerHandler{
		options: handler.Options{
			RateLimiter: &mockRateLimiterContainer{
				limiterFn: func(key string) rate.Limiter {
					return &mockRateLimiter{
						allowFn: func(n int) bool { return true },
					}
				},
			},
		},
	}
	addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}
	if !h.checkRateLimit(addr) {
		t.Error("checkRateLimit returned false, want true")
	}
}

func TestCheckRateLimit_Denied(t *testing.T) {
	h := &routerHandler{
		options: handler.Options{
			RateLimiter: &mockRateLimiterContainer{
				limiterFn: func(key string) rate.Limiter {
					return &mockRateLimiter{
						allowFn: func(n int) bool { return false },
					}
				},
			},
		},
	}
	addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}
	if h.checkRateLimit(addr) {
		t.Error("checkRateLimit returned true, want false")
	}
}

func TestCheckRateLimit_LimiterByHost(t *testing.T) {
	var gotKey string
	h := &routerHandler{
		options: handler.Options{
			RateLimiter: &mockRateLimiterContainer{
				limiterFn: func(key string) rate.Limiter {
					gotKey = key
					return &mockRateLimiter{
						allowFn: func(n int) bool { return true },
					}
				},
			},
		},
	}
	addr := &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 12345}
	h.checkRateLimit(addr)
	if gotKey != "10.0.0.1" {
		t.Errorf("limiter key = %q, want 10.0.0.1", gotKey)
	}
}

// ---------------------------------------------------------------------------
// observeStats tests
// ---------------------------------------------------------------------------

// newHandlerWithObserver creates a routerHandler with a fake observer.
func newHandlerWithObserver(t *testing.T, obs *fakeObserver) *routerHandler {
	t.Helper()
	h := &routerHandler{
		options: handler.Options{
			Observer: obs,
			Service:  "test-svc",
		},
		md: metadata{
			observerPeriod:       50 * time.Millisecond,
			observerResetTraffic: false,
		},
	}
	h.stats = stats_util.NewHandlerStats("test-svc", false)
	return h
}

func TestObserveStats_NilObserver(t *testing.T) {
	h := &routerHandler{}
	// Should return immediately without panic
	h.observeStats(context.Background())
}

func TestObserveStats_NormalCycle(t *testing.T) {
	obs := newFakeObserver(10)
	h := newHandlerWithObserver(t, obs)

	// Simulate some stat activity so Events() returns data.
	h.stats.Stats("client-1").Add(stats.KindTotalConns, 1)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		h.observeStats(ctx)
	}()

	// Wait for at least one observation cycle
	select {
	case <-obs.Events():
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for observer event")
	}

	cancel()
	wg.Wait()
}

func TestObserveStats_ContextCancel(t *testing.T) {
	obs := newFakeObserver(10)
	h := newHandlerWithObserver(t, obs)

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel immediately

	// Should return immediately when context is already cancelled
	done := make(chan struct{})
	go func() {
		h.observeStats(ctx)
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("observeStats did not exit after context cancel")
	}
}

func TestObserveStats_RetryOnError(t *testing.T) {
	var callCount int
	var mu sync.Mutex
	obs := &fakeObserver{
		eventsCh: make(chan []observer.Event, 10),
		errFunc: func() error {
			mu.Lock()
			callCount++
			count := callCount
			mu.Unlock()
			if count <= 2 {
				return errors.New("observer error")
			}
			return nil
		},
	}
	h := newHandlerWithObserver(t, obs)

	// Simulate stat activity so Events() returns data.
	h.stats.Stats("client-1").Add(stats.KindTotalConns, 1)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		h.observeStats(ctx)
	}()

	// Wait for successful observation (callCount > 2)
	select {
	case <-obs.Events():
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for successful observer event")
	}

	mu.Lock()
	if callCount < 2 {
		t.Errorf("callCount = %d, want at least 2 (error retries)", callCount)
	}
	mu.Unlock()

	cancel()
	wg.Wait()
}