Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/ops_ingress_reject_test.go

@@ -0,0 +1,115 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type ingressRejectRepoStub struct {
+	mu        sync.Mutex
+	failCount int
+	calls     int
+	requests  int64
+}
+
+func (r *ingressRejectRepoStub) BatchUpsertIngressRejects(_ context.Context, items []*OpsIngressRejectAggregate) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.calls++
+	if r.failCount > 0 {
+		r.failCount--
+		return errors.New("temporary database failure")
+	}
+	for _, item := range items {
+		if item != nil {
+			r.requests += item.RequestCount
+		}
+	}
+	return nil
+}
+
+func (r *ingressRejectRepoStub) ListIngressRejects(context.Context, *OpsIngressRejectFilter) (*OpsIngressRejectList, error) {
+	return &OpsIngressRejectList{}, nil
+}
+
+func TestOpsIngressRejectAggregatorUsesGlobalBucketCapacity(t *testing.T) {
+	repo := &ingressRejectRepoStub{}
+	a := NewOpsIngressRejectAggregator(repo)
+	a.Start()
+
+	// Concentrating all dimensions in one shard must not waste capacity in the others.
+	inserted := 0
+	for i := 0; inserted < 600; i++ {
+		ip := fmt.Sprintf("target-%d", i)
+		key := ingressRejectKey{reason: "invalid_api_key", routeFamily: "messages", protocol: "anthropic", clientIP: ip}
+		if ingressRejectHash(key)%ingressRejectShardCount == 0 {
+			a.RecordIngressReject(key.reason, key.routeFamily, key.protocol, ip, 0, 0)
+			inserted++
+		}
+	}
+	require.Equal(t, int64(600), a.Health().Cardinality)
+
+	for i := 0; i < ingressRejectMaxEntries; i++ {
+		a.RecordIngressReject("invalid_api_key", "messages", "anthropic", fmt.Sprintf("rotating-%d", i), 0, 0)
+	}
+	health := a.Health()
+	require.Equal(t, int64(ingressRejectMaxEntries), health.Cardinality)
+	require.Greater(t, health.Overflowed, uint64(0))
+
+	a.snapshotAndEnqueue(false)
+	require.Equal(t, int64(ingressRejectMaxEntries), a.Health().Cardinality, "periodic flush must retain the minute budget")
+	a.Stop()
+}
+
+func TestOpsIngressRejectAggregatorConcurrentCountAndStopFlush(t *testing.T) {
+	repo := &ingressRejectRepoStub{}
+	a := NewOpsIngressRejectAggregator(repo)
+	a.Start()
+
+	const goroutines = 32
+	const perGoroutine = 200
+	var wg sync.WaitGroup
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < perGoroutine; j++ {
+				a.RecordIngressReject("invalid_api_key", "responses", "openai", "192.0.2.10", 0, 0)
+			}
+		}()
+	}
+	wg.Wait()
+	a.Stop()
+
+	repo.mu.Lock()
+	require.Equal(t, int64(goroutines*perGoroutine), repo.requests)
+	repo.mu.Unlock()
+	require.False(t, a.Health().Accepting)
+	a.RecordIngressReject("invalid_api_key", "responses", "openai", "192.0.2.10", 0, 0)
+}
+
+func TestOpsIngressRejectAggregatorRetriesBoundedPendingBatch(t *testing.T) {
+	repo := &ingressRejectRepoStub{failCount: 1}
+	a := NewOpsIngressRejectAggregator(repo)
+	a.accepting.Store(true)
+	a.RecordIngressReject("group_deleted", "messages", "anthropic", "192.0.2.20", 1, 2)
+	a.snapshotAndEnqueue(false)
+	a.flushPending()
+	health := a.Health()
+	require.Equal(t, 1, health.PendingBatches)
+	require.Equal(t, uint64(1), health.FlushFailures)
+	require.Equal(t, uint64(0), health.Dropped)
+
+	a.flushPending()
+	require.Equal(t, 0, a.Health().PendingBatches)
+	a.Stop()
+	repo.mu.Lock()
+	require.Equal(t, int64(1), repo.requests)
+	require.GreaterOrEqual(t, repo.calls, 2)
+	repo.mu.Unlock()
+}