25'ten fazla konu seçemezsiniz Konular bir harf veya rakamla başlamalı, kısa çizgiler ('-') içerebilir ve en fazla 35 karakter uzunluğunda olabilir.
 
 
 

243 satır
7.5 KiB

  1. package metrics
  2. import (
  3. "net/http/httptest"
  4. "testing"
  5. "time"
  6. "github.com/gin-gonic/gin"
  7. "github.com/prometheus/client_golang/prometheus"
  8. "github.com/prometheus/client_golang/prometheus/testutil"
  9. "github.com/stretchr/testify/assert"
  10. )
  11. func newTestRegistry() *prometheus.Registry {
  12. return prometheus.NewRegistry()
  13. }
  14. func TestRequestsTotalCounter(t *testing.T) {
  15. reg := newTestRegistry()
  16. m := NewMetrics(reg)
  17. m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
  18. m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
  19. m.RequestsTotal.WithLabelValues("2", "claude-3", "429", "false", "chat").Inc()
  20. assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat")))
  21. assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("2", "claude-3", "429", "false", "chat")))
  22. }
  23. func TestRequestErrorsTotalCounter(t *testing.T) {
  24. reg := newTestRegistry()
  25. m := NewMetrics(reg)
  26. m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429").Inc()
  27. m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429").Inc()
  28. m.RequestErrorsTotal.WithLabelValues("2", "claude-3", "timeout", "context_deadline_exceeded").Inc()
  29. assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestErrorsTotal.WithLabelValues("1", "gpt-4", "upstream", "429")))
  30. assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestErrorsTotal.WithLabelValues("2", "claude-3", "timeout", "context_deadline_exceeded")))
  31. }
  32. func TestRequestRetriesTotalCounter(t *testing.T) {
  33. reg := newTestRegistry()
  34. m := NewMetrics(reg)
  35. m.RequestRetriesTotal.WithLabelValues("1", "gpt-4").Inc()
  36. m.RequestRetriesTotal.WithLabelValues("1", "gpt-4").Inc()
  37. assert.Equal(t, float64(2), testutil.ToFloat64(m.RequestRetriesTotal.WithLabelValues("1", "gpt-4")))
  38. }
  39. func TestTokensTotalCounter(t *testing.T) {
  40. reg := newTestRegistry()
  41. m := NewMetrics(reg)
  42. m.TokensTotal.WithLabelValues("1", "gpt-4", "prompt").Add(100)
  43. m.TokensTotal.WithLabelValues("1", "gpt-4", "completion").Add(50)
  44. m.TokensTotal.WithLabelValues("1", "gpt-4", "cache").Add(20)
  45. assert.Equal(t, float64(100), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "prompt")))
  46. assert.Equal(t, float64(50), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "completion")))
  47. assert.Equal(t, float64(20), testutil.ToFloat64(m.TokensTotal.WithLabelValues("1", "gpt-4", "cache")))
  48. }
  49. func TestQuotaConsumedTotalCounter(t *testing.T) {
  50. reg := newTestRegistry()
  51. m := NewMetrics(reg)
  52. m.QuotaConsumedTotal.WithLabelValues("1", "gpt-4", "wallet").Add(5000)
  53. assert.Equal(t, float64(5000), testutil.ToFloat64(m.QuotaConsumedTotal.WithLabelValues("1", "gpt-4", "wallet")))
  54. }
  55. func TestRequestDurationHistogram(t *testing.T) {
  56. reg := newTestRegistry()
  57. m := NewMetrics(reg)
  58. m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(1.5)
  59. m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(2.5)
  60. // 验证 histogram metric family 被注册且有数据
  61. assert.Equal(t, 1, testutil.CollectAndCount(m.RequestDuration))
  62. }
  63. func TestUpstreamDurationHistogram(t *testing.T) {
  64. reg := newTestRegistry()
  65. m := NewMetrics(reg)
  66. m.UpstreamDuration.WithLabelValues("1", "gpt-4").Observe(0.5)
  67. assert.Equal(t, 1, testutil.CollectAndCount(m.UpstreamDuration))
  68. }
  69. func TestFirstTokenDurationHistogram(t *testing.T) {
  70. reg := newTestRegistry()
  71. m := NewMetrics(reg)
  72. m.FirstTokenDuration.WithLabelValues("1", "gpt-4").Observe(0.1)
  73. assert.Equal(t, 1, testutil.CollectAndCount(m.FirstTokenDuration))
  74. }
  75. func TestActiveRequestsGauge(t *testing.T) {
  76. reg := newTestRegistry()
  77. m := NewMetrics(reg)
  78. m.ActiveRequests.WithLabelValues("gpt-4").Inc()
  79. m.ActiveRequests.WithLabelValues("gpt-4").Inc()
  80. m.ActiveRequests.WithLabelValues("claude-3").Inc()
  81. assert.Equal(t, float64(2), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("gpt-4")))
  82. assert.Equal(t, float64(1), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("claude-3")))
  83. m.ActiveRequests.WithLabelValues("gpt-4").Dec()
  84. assert.Equal(t, float64(1), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("gpt-4")))
  85. }
  86. func TestMetricsDisabledIsNoOp(t *testing.T) {
  87. // 当 METRICS_ENABLED=false 时,使用独立 registry(Discard)
  88. reg := prometheus.NewPedanticRegistry()
  89. m := NewMetrics(reg)
  90. // 这些调用不应该 panic
  91. m.RequestsTotal.WithLabelValues("1", "gpt-4", "200", "true", "chat").Inc()
  92. m.ActiveRequests.WithLabelValues("gpt-4").Inc()
  93. m.RequestDuration.WithLabelValues("1", "gpt-4", "true", "chat").Observe(1.0)
  94. }
  95. func TestMetricsMiddleware_ActiveRequests(t *testing.T) {
  96. gin.SetMode(gin.TestMode)
  97. reg := newTestRegistry()
  98. m := NewMetrics(reg)
  99. r := gin.New()
  100. r.Use(MetricsMiddleware(m))
  101. blockCh := make(chan struct{})
  102. r.GET("/test", func(c *gin.Context) {
  103. <-blockCh
  104. c.Status(200)
  105. })
  106. go func() {
  107. req := httptest.NewRequest("GET", "/test", nil)
  108. w := httptest.NewRecorder()
  109. r.ServeHTTP(w, req)
  110. }()
  111. go func() {
  112. req := httptest.NewRequest("GET", "/test", nil)
  113. w := httptest.NewRecorder()
  114. r.ServeHTTP(w, req)
  115. }()
  116. time.Sleep(100 * time.Millisecond)
  117. assert.Equal(t, float64(2), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("")))
  118. close(blockCh)
  119. time.Sleep(100 * time.Millisecond)
  120. assert.Equal(t, float64(0), testutil.ToFloat64(m.ActiveRequests.WithLabelValues("")))
  121. }
  122. func TestMetricsMiddleware_RequestsTotalByCode(t *testing.T) {
  123. gin.SetMode(gin.TestMode)
  124. reg := newTestRegistry()
  125. m := NewMetrics(reg)
  126. r := gin.New()
  127. r.Use(MetricsMiddleware(m))
  128. r.GET("/200", func(c *gin.Context) { c.Status(200) })
  129. r.GET("/400", func(c *gin.Context) { c.Status(400) })
  130. r.GET("/500", func(c *gin.Context) { c.Status(500) })
  131. for i := 0; i < 3; i++ {
  132. req := httptest.NewRequest("GET", "/200", nil)
  133. w := httptest.NewRecorder()
  134. r.ServeHTTP(w, req)
  135. }
  136. req := httptest.NewRequest("GET", "/400", nil)
  137. w := httptest.NewRecorder()
  138. r.ServeHTTP(w, req)
  139. req = httptest.NewRequest("GET", "/500", nil)
  140. w = httptest.NewRecorder()
  141. r.ServeHTTP(w, req)
  142. assert.Equal(t, float64(3), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "200", "", "")))
  143. assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "400", "", "")))
  144. assert.Equal(t, float64(1), testutil.ToFloat64(m.RequestsTotal.WithLabelValues("", "", "500", "", "")))
  145. }
  146. func TestMetricsMiddleware_RequestDuration(t *testing.T) {
  147. gin.SetMode(gin.TestMode)
  148. reg := newTestRegistry()
  149. m := NewMetrics(reg)
  150. r := gin.New()
  151. r.Use(MetricsMiddleware(m))
  152. r.GET("/slow", func(c *gin.Context) {
  153. time.Sleep(100 * time.Millisecond)
  154. c.Status(200)
  155. })
  156. req := httptest.NewRequest("GET", "/slow", nil)
  157. w := httptest.NewRecorder()
  158. r.ServeHTTP(w, req)
  159. assert.Equal(t, 1, testutil.CollectAndCount(m.RequestDuration))
  160. }
  161. func TestMetricsEndpoint_Format(t *testing.T) {
  162. handler := NewMetricsHandler()
  163. req := httptest.NewRequest("GET", "/metrics", nil)
  164. w := httptest.NewRecorder()
  165. handler.ServeHTTP(w, req)
  166. assert.Equal(t, 200, w.Code)
  167. assert.Contains(t, w.Header().Get("Content-Type"), "text/plain")
  168. }
  169. func TestMetricsEndpoint_WithData(t *testing.T) {
  170. reg := newTestRegistry()
  171. _ = NewMetrics(reg)
  172. handler := NewMetricsHandlerWithRegistry(reg)
  173. req := httptest.NewRequest("GET", "/metrics", nil)
  174. w := httptest.NewRecorder()
  175. handler.ServeHTTP(w, req)
  176. assert.Equal(t, 200, w.Code)
  177. }
  178. func TestMetricsEndpoint_EmptyNoPanic(t *testing.T) {
  179. handler := NewMetricsHandler()
  180. req := httptest.NewRequest("GET", "/metrics", nil)
  181. w := httptest.NewRecorder()
  182. handler.ServeHTTP(w, req)
  183. assert.Equal(t, 200, w.Code)
  184. }