new-api/controller/model_meta_test.go

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