|
- package metrics
-
- import (
- "testing"
-
- "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)
- }
|