sub2api/backend/internal/server/middleware/ingress_reject_access_sampler_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestIngressRejectAccessSamplerConcurrentGlobalLimit(t *testing.T) {
	sampler := newIngressRejectAccessSampler(10, time.Hour, time.Minute)
	now := time.Now()
	var allowed atomic.Int64
	var wg sync.WaitGroup
	for i := 0; i < 200; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if ok, _ := sampler.allow(now); ok {
				allowed.Add(1)
			}
		}()
	}
	wg.Wait()
	require.Equal(t, int64(10), allowed.Load())
}

func TestLoggerIngressRejectSamplingIsBoundedAndSummarySkipsOpsSink(t *testing.T) {
	gin.SetMode(gin.TestMode)
	original := globalIngressRejectAccessSampler
	globalIngressRejectAccessSampler = newIngressRejectAccessSampler(2, time.Hour, time.Hour)
	t.Cleanup(func() { globalIngressRejectAccessSampler = original })
	sink := initMiddlewareTestLogger(t)
	router := gin.New()
	router.Use(Logger())
	router.GET("/v1/messages", func(c *gin.Context) {
		MarkIngressRejected(c, IngressRejectInvalidAPIKey)
		c.Status(http.StatusUnauthorized)
	})
	for i := 0; i < 20; i++ {
		router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/messages", nil))
	}
	var accessEvents, summaries int
	for _, event := range sink.list() {
		switch event.Message {
		case "http request completed":
			accessEvents++
		case "ingress rejection access logs dropped":
			summaries++
			if skipped, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); !skipped {
				t.Fatalf("dropped summary must skip ops system log sink")
			}
		}
	}
	require.Equal(t, 2, accessEvents)
	require.Equal(t, 1, summaries)
}

func TestIngressRejectAccessSamplerDroppedSummaryIsLowFrequency(t *testing.T) {
	sampler := newIngressRejectAccessSampler(1, time.Hour, time.Second)
	now := time.Now()
	allowed, summary := sampler.allow(now)
	require.True(t, allowed)
	require.Zero(t, summary)

	allowed, summary = sampler.allow(now.Add(100 * time.Millisecond))
	require.False(t, allowed)
	require.Equal(t, uint64(1), summary)
	allowed, summary = sampler.allow(now.Add(200 * time.Millisecond))
	require.False(t, allowed)
	require.Zero(t, summary)
	allowed, summary = sampler.allow(now.Add(2 * time.Second))
	require.False(t, allowed)
	require.Equal(t, uint64(2), summary)
}