|
|
|
@@ -0,0 +1,93 @@ |
|
|
|
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") |
|
|
|
} |
|
|
|
} |