package controller import ( "bytes" "net/http" "net/http/httptest" "strconv" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupModelMetaTestDB(t *testing.T) { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) require.NoError(t, db.AutoMigrate(&model.Model{})) origDB := model.DB model.DB = db t.Cleanup(func() { model.RefreshPricing() model.DB = origDB }) } func TestCreateModelMetaDefaultsEmptyEndpointsToOpenAI(t *testing.T) { setupModelMetaTestDB(t) gin.SetMode(gin.TestMode) r := gin.New() r.POST("/api/models/", CreateModelMeta) req := httptest.NewRequest( http.MethodPost, "/api/models/", bytes.NewBufferString(`{"model_name":"custom-chat-model"}`), ) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() r.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) var resp map[string]any require.NoError(t, common.Unmarshal(w.Body.Bytes(), &resp)) require.True(t, resp["success"].(bool)) var saved model.Model require.NoError(t, model.DB.Where("model_name = ?", "custom-chat-model").First(&saved).Error) require.Equal(t, `["openai"]`, saved.Endpoints) } func TestUpdateModelMetaDefaultsEmptyEndpointsToOpenAI(t *testing.T) { setupModelMetaTestDB(t) gin.SetMode(gin.TestMode) existing := &model.Model{ ModelName: "custom-chat-model", Endpoints: `{"anthropic":{"path":"/v1/messages","method":"POST"}}`, Status: 1, } require.NoError(t, existing.Insert()) r := gin.New() r.PUT("/api/models/", UpdateModelMeta) req := httptest.NewRequest( http.MethodPut, "/api/models/", bytes.NewBufferString(`{"id":`+strconv.Itoa(existing.Id)+`,"model_name":"custom-chat-model","endpoints":"","status":1}`), ) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() r.ServeHTTP(w, req) require.Equal(t, http.StatusOK, w.Code) var resp map[string]any require.NoError(t, common.Unmarshal(w.Body.Bytes(), &resp)) require.True(t, resp["success"].(bool)) var saved model.Model require.NoError(t, model.DB.First(&saved, existing.Id).Error) require.Equal(t, `["openai"]`, saved.Endpoints) }