|
- package metrics
-
- import (
- "testing"
-
- "github.com/glebarez/sqlite"
- "github.com/prometheus/client_golang/prometheus"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- )
-
- func setupCollectorTestDB(t *testing.T) *gorm.DB {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- model.DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&model.Channel{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- sqlDB.Close()
- })
-
- return db
- }
-
- func TestChannelCollector_CollectsChannelStatus(t *testing.T) {
- db := setupCollectorTestDB(t)
-
- ch1 := model.Channel{Id: 1, Name: "openai-1", Type: 1, Status: common.ChannelStatusEnabled}
- ch2 := model.Channel{Id: 2, Name: "claude-1", Type: 14, Status: common.ChannelStatusEnabled}
- ch3 := model.Channel{Id: 3, Name: "gemini-1", Type: 24, Status: 3}
- require.NoError(t, db.Create(&ch1).Error)
- require.NoError(t, db.Create(&ch2).Error)
- require.NoError(t, db.Create(&ch3).Error)
-
- collector := NewChannelCollector()
- reg := prometheus.NewRegistry()
- reg.MustRegister(collector)
-
- mfs, err := reg.Gather()
- require.NoError(t, err)
-
- var found bool
- for _, mf := range mfs {
- if mf.GetName() == "newapi_channel_status" {
- found = true
- metrics := mf.GetMetric()
- assert.Len(t, metrics, 3)
-
- statusMap := map[string]float64{}
- for _, m := range metrics {
- chId := ""
- for _, label := range m.GetLabel() {
- if label.GetName() == "channel" {
- chId = label.GetValue()
- }
- }
- statusMap[chId] = m.GetGauge().GetValue()
- }
- assert.Equal(t, float64(1), statusMap["1"])
- assert.Equal(t, float64(1), statusMap["2"])
- assert.Equal(t, float64(0), statusMap["3"])
- }
- }
- assert.True(t, found, "newapi_channel_status metric not found")
- }
-
- func TestChannelCollector_EmptyDB(t *testing.T) {
- setupCollectorTestDB(t)
-
- collector := NewChannelCollector()
- reg := prometheus.NewRegistry()
- reg.MustRegister(collector)
-
- mfs, err := reg.Gather()
- require.NoError(t, err)
-
- for _, mf := range mfs {
- assert.NotEqual(t, "newapi_channel_status", mf.GetName(), "should not have channel status with empty DB")
- }
- }
|