Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

94 righe
2.3 KiB

  1. package metrics
  2. import (
  3. "testing"
  4. "github.com/glebarez/sqlite"
  5. "github.com/prometheus/client_golang/prometheus"
  6. "github.com/stretchr/testify/assert"
  7. "github.com/stretchr/testify/require"
  8. "gorm.io/gorm"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/model"
  11. )
  12. func setupCollectorTestDB(t *testing.T) *gorm.DB {
  13. t.Helper()
  14. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  15. require.NoError(t, err)
  16. sqlDB, _ := db.DB()
  17. sqlDB.SetMaxOpenConns(1)
  18. origDB := model.DB
  19. model.DB = db
  20. common.UsingSQLite = true
  21. common.RedisEnabled = false
  22. require.NoError(t, db.AutoMigrate(&model.Channel{}))
  23. t.Cleanup(func() {
  24. model.DB = origDB
  25. sqlDB.Close()
  26. })
  27. return db
  28. }
  29. func TestChannelCollector_CollectsChannelStatus(t *testing.T) {
  30. db := setupCollectorTestDB(t)
  31. ch1 := model.Channel{Id: 1, Name: "openai-1", Type: 1, Status: common.ChannelStatusEnabled}
  32. ch2 := model.Channel{Id: 2, Name: "claude-1", Type: 14, Status: common.ChannelStatusEnabled}
  33. ch3 := model.Channel{Id: 3, Name: "gemini-1", Type: 24, Status: 3}
  34. require.NoError(t, db.Create(&ch1).Error)
  35. require.NoError(t, db.Create(&ch2).Error)
  36. require.NoError(t, db.Create(&ch3).Error)
  37. collector := NewChannelCollector()
  38. reg := prometheus.NewRegistry()
  39. reg.MustRegister(collector)
  40. mfs, err := reg.Gather()
  41. require.NoError(t, err)
  42. var found bool
  43. for _, mf := range mfs {
  44. if mf.GetName() == "newapi_channel_status" {
  45. found = true
  46. metrics := mf.GetMetric()
  47. assert.Len(t, metrics, 3)
  48. statusMap := map[string]float64{}
  49. for _, m := range metrics {
  50. chId := ""
  51. for _, label := range m.GetLabel() {
  52. if label.GetName() == "channel" {
  53. chId = label.GetValue()
  54. }
  55. }
  56. statusMap[chId] = m.GetGauge().GetValue()
  57. }
  58. assert.Equal(t, float64(1), statusMap["1"])
  59. assert.Equal(t, float64(1), statusMap["2"])
  60. assert.Equal(t, float64(0), statusMap["3"])
  61. }
  62. }
  63. assert.True(t, found, "newapi_channel_status metric not found")
  64. }
  65. func TestChannelCollector_EmptyDB(t *testing.T) {
  66. setupCollectorTestDB(t)
  67. collector := NewChannelCollector()
  68. reg := prometheus.NewRegistry()
  69. reg.MustRegister(collector)
  70. mfs, err := reg.Gather()
  71. require.NoError(t, err)
  72. for _, mf := range mfs {
  73. assert.NotEqual(t, "newapi_channel_status", mf.GetName(), "should not have channel status with empty DB")
  74. }
  75. }