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") } }