|
- package metrics
-
- import (
- "net/http/httptest"
- "testing"
- "time"
-
- "github.com/gin-gonic/gin"
- "github.com/prometheus/client_golang/prometheus"
- "github.com/prometheus/client_golang/prometheus/testutil"
- "github.com/stretchr/testify/assert"
- )
-
- func newTestRegistry() *prometheus.Registry {
- return prometheus.NewRegistry()
- }
-
- func TestRequestsTotalCounter(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
- m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
- m.RequestsTotal.WithLabelValues("2", "claude-3", "429", "false", "chat").Inc()
-
- assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat")))
- assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("2", "claude-3", "429", "false", "chat")))
- }
-
- func TestRequestErrorsTotalCounter(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429").Inc()
- m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429").Inc()
- m.RequestErrorsTotal.WithLabelValues("2", "claude-3", "timeout", "context_deadline_exceeded").Inc()
-
- assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429")))
- assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestErrorsTotal.WithLabelValues("2", "claude-3", "timeout", "context_deadline_exceeded")))
- }
-
- func TestRequestRetriesTotalCounter(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.RequestRetriesTotal.WithLabelValues("1", "gpt-4").Inc()
- m.RequestRetriesTotal.WithLabelValues("1", "gpt-4").Inc()
-
- assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestRetriesTotal.WithLabelValues("1", "gpt-4")))
- }
-
- func TestTokensTotalCounter(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.TokensTotal.WithLabelValues("1", "gpt-4", "prompt").Add(100)
- m.TokensTotal.WithLabelValues("1", "gpt-4", "completion").Add(50)
- m.TokensTotal.WithLabelValues("1", "gpt-4", "cache").Add(20)
-
- assert.Equal(t, float64(100), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "prompt")))
- assert.Equal(t, float64(50), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "completion")))
- assert.Equal(t, float64(20), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "cache")))
- }
-
- func TestQuotaConsumedTotalCounter(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.QuotaConsumedTotal.WithLabelValues("1", "gpt-4", "wallet").Add(5000)
-
- assert.Equal(t, float64(5000), testutil.ToFloat64(m.QuotaConsumedTotal.WithLabelValues("1", "gpt-4", "wallet")))
- }
-
- func TestRequestDurationHistogram(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(1.5)
- m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(2.5)
-
- // 验证 histogram metric family 被注册且有数据
- assert.Equal(t, 1, testutil.CollectAndCount(m.RequestDuration))
- }
-
- func TestUpstreamDurationHistogram(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.UpstreamDuration.WithLabelValues("1", "gpt-4").Observe(0.5)
-
- assert.Equal(t, 1, testutil.CollectAndCount(m.UpstreamDuration))
- }
-
- func TestFirstTokenDurationHistogram(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.FirstTokenDuration.WithLabelValues("1", "gpt-4").Observe(0.1)
-
- assert.Equal(t, 1, testutil.CollectAndCount(m.FirstTokenDuration))
- }
-
- func TestActiveRequestsGauge(t *testing.T) {
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- m.ActiveRequests.WithLabelValues("gpt-4").Inc()
- m.ActiveRequests.WithLabelValues("gpt-4").Inc()
- m.ActiveRequests.WithLabelValues("claude-3").Inc()
-
- assert.Equal(t, float64(2), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("gpt-4")))
- assert.Equal(t, float64(1), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("claude-3")))
-
- m.ActiveRequests.WithLabelValues("gpt-4").Dec()
- assert.Equal(t, float64(1), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("gpt-4")))
- }
-
- func TestMetricsDisabledIsNoOp(t *testing.T) {
- // 当 METRICS_ENABLED=false 时,使用独立 registry(Discard)
- reg := prometheus.NewPedanticRegistry()
- m := NewMetrics(reg)
- // 这些调用不应该 panic
- m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
- m.ActiveRequests.WithLabelValues("gpt-4").Inc()
- m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(1.0)
- }
-
- func TestMetricsMiddleware_ActiveRequests(t *testing.T) {
- gin.SetMode(gin.TestMode)
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- r := gin.New()
- r.Use(MetricsMiddleware(m))
-
- blockCh := make(chan struct{})
- r.GET("/test", func(c *gin.Context) {
- <-blockCh
- c.Status(200)
- })
-
- go func() {
- req := httptest.NewRequest("GET", "/test", nil)
- w := httptest.NewRecorder()
- r.ServeHTTP(w, req)
- }()
- go func() {
- req := httptest.NewRequest("GET", "/test", nil)
- w := httptest.NewRecorder()
- r.ServeHTTP(w, req)
- }()
-
- time.Sleep(100 * time.Millisecond)
-
- assert.Equal(t, float64(2), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("")))
-
- close(blockCh)
- time.Sleep(100 * time.Millisecond)
-
- assert.Equal(t, float64(0), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("")))
- }
-
- func TestMetricsMiddleware_RequestsTotalByCode(t *testing.T) {
- gin.SetMode(gin.TestMode)
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- r := gin.New()
- r.Use(MetricsMiddleware(m))
- r.GET("/200", func(c *gin.Context) { c.Status(200) })
- r.GET("/400", func(c *gin.Context) { c.Status(400) })
- r.GET("/500", func(c *gin.Context) { c.Status(500) })
-
- for i := 0; i < 3; i++ {
- req := httptest.NewRequest("GET", "/200", nil)
- w := httptest.NewRecorder()
- r.ServeHTTP(w, req)
- }
- req := httptest.NewRequest("GET", "/400", nil)
- w := httptest.NewRecorder()
- r.ServeHTTP(w, req)
- req = httptest.NewRequest("GET", "/500", nil)
- w = httptest.NewRecorder()
- r.ServeHTTP(w, req)
-
- assert.Equal(t, float64(3), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "200", "", "")))
- assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "400", "", "")))
- assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "500", "", "")))
- }
-
- func TestMetricsMiddleware_RequestDuration(t *testing.T) {
- gin.SetMode(gin.TestMode)
- reg := newTestRegistry()
- m := NewMetrics(reg)
-
- r := gin.New()
- r.Use(MetricsMiddleware(m))
- r.GET("/slow", func(c *gin.Context) {
- time.Sleep(100 * time.Millisecond)
- c.Status(200)
- })
-
- req := httptest.NewRequest("GET", "/slow", nil)
- w := httptest.NewRecorder()
- r.ServeHTTP(w, req)
-
- assert.Equal(t, 1, testutil.CollectAndCount(m.RequestDuration))
- }
-
- func TestMetricsEndpoint_Format(t *testing.T) {
- handler := NewMetricsHandler()
-
- req := httptest.NewRequest("GET", "/metrics", nil)
- w := httptest.NewRecorder()
- handler.ServeHTTP(w, req)
-
- assert.Equal(t, 200, w.Code)
- assert.Contains(t, w.Header().Get("Content-Type"), "text/plain")
- }
-
- func TestMetricsEndpoint_WithData(t *testing.T) {
- reg := newTestRegistry()
- _ = NewMetrics(reg)
-
- handler := NewMetricsHandlerWithRegistry(reg)
-
- req := httptest.NewRequest("GET", "/metrics", nil)
- w := httptest.NewRecorder()
- handler.ServeHTTP(w, req)
-
- assert.Equal(t, 200, w.Code)
- }
-
- func TestMetricsEndpoint_EmptyNoPanic(t *testing.T) {
- handler := NewMetricsHandler()
-
- req := httptest.NewRequest("GET", "/metrics", nil)
- w := httptest.NewRecorder()
- handler.ServeHTTP(w, req)
-
- assert.Equal(t, 200, w.Code)
- }
|