new-api/service/channel_monitor_checker_test.go

package service

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupChannelMonitorTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false

	previousDB := model.DB
	previousLogDB := model.LOG_DB
	dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&model.ChannelMonitor{}, &model.ChannelMonitorHistory{}, &model.ChannelMonitorDailyRollup{}))

	model.DB = db
	model.LOG_DB = db

	t.Cleanup(func() {
		model.DB = previousDB
		model.LOG_DB = previousLogDB
		sqlDB, err := db.DB()
		if err == nil {
			_ = sqlDB.Close()
		}
	})
	return db
}

func allowTestChannelMonitorURL(string) error {
	return nil
}

func TestRunMonitorCheckPassWithCustomTemplateBodyAndHeaders(t *testing.T) {
	db := setupChannelMonitorTestDB(t)

	var receivedBody string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "ok", r.Header.Get("X-Test"))
		body, err := io.ReadAll(r.Body)
		require.NoError(t, err)
		receivedBody = string(body)
		_, _ = w.Write([]byte(`{"id":"chatcmpl_test"}`))
	}))
	defer server.Close()

	monitor := &model.ChannelMonitor{
		Name:             "custom",
		ChannelId:        7,
		TargetUrl:        server.URL,
		PrimaryModel:     "m-custom",
		RequestTemplates: `["custom prompt"]`,
		CustomBody:       `{"model":"{{model}}","input":"{{template}}"}`,
		CustomHeaders:    `{"X-Test":"ok"}`,
		Timeout:          5,
		Status:           1,
	}
	require.NoError(t, db.Create(monitor).Error)

	err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
		client:      server.Client(),
		validateURL: allowTestChannelMonitorURL,
	})

	require.NoError(t, err)
	require.Contains(t, receivedBody, `"model":"m-custom"`)
	require.Contains(t, receivedBody, `"input":"custom prompt"`)
	require.NotNil(t, monitor.LastCheckedAt)

	var history model.ChannelMonitorHistory
	require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
	require.Equal(t, monitor.Id, history.MonitorId)
	require.Equal(t, 7, history.ChannelId)
	require.Equal(t, "m-custom", history.ModelName)
	require.Equal(t, "pass", history.Status)
	require.Equal(t, 200, history.StatusCode)
	require.GreaterOrEqual(t, history.LatencyMs, 0)
	require.Empty(t, history.ErrorMessage)
	require.Contains(t, history.RequestBody, `"model":"m-custom"`)
	require.Contains(t, history.ResponseSnippet, "chatcmpl_test")
	require.False(t, history.CheckedAt.IsZero())
}

func TestRunMonitorCheckRecordsFailedStatusCode(t *testing.T) {
	db := setupChannelMonitorTestDB(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write([]byte(`{"error":"bad"}`))
	}))
	defer server.Close()

	monitor := &model.ChannelMonitor{
		Name:         "failed",
		TargetUrl:    server.URL,
		PrimaryModel: "m-failed",
		Timeout:      5,
		Status:       1,
	}
	require.NoError(t, db.Create(monitor).Error)

	err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
		client:      server.Client(),
		validateURL: allowTestChannelMonitorURL,
	})

	require.NoError(t, err)
	var history model.ChannelMonitorHistory
	require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
	require.Equal(t, "failed", history.Status)
	require.Equal(t, http.StatusInternalServerError, history.StatusCode)
	require.Contains(t, history.ResponseSnippet, "bad")
}

func TestRunMonitorCheckRecordsSSRFFailure(t *testing.T) {
	db := setupChannelMonitorTestDB(t)

	monitor := &model.ChannelMonitor{
		Name:         "ssrf",
		TargetUrl:    "http://127.0.0.1:1/v1/chat/completions",
		PrimaryModel: "m-ssrf",
		Timeout:      5,
		Status:       1,
	}
	require.NoError(t, db.Create(monitor).Error)

	err := RunMonitorCheck(monitor)

	require.Error(t, err)
	var history model.ChannelMonitorHistory
	require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
	require.Equal(t, "error", history.Status)
	require.Contains(t, history.ErrorMessage, "SSRF validation failed")
	require.Contains(t, history.RequestBody, "m-ssrf")
	require.Equal(t, 0, history.StatusCode)
}

func TestRunMonitorCheckRecordsInvalidCustomHeaders(t *testing.T) {
	db := setupChannelMonitorTestDB(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatal("request should not be sent when custom headers are invalid")
	}))
	defer server.Close()

	monitor := &model.ChannelMonitor{
		Name:          "headers",
		TargetUrl:     server.URL,
		PrimaryModel:  "m-headers",
		CustomHeaders: `{"X-Test":123}`,
		Timeout:       5,
		Status:        1,
	}
	require.NoError(t, db.Create(monitor).Error)

	err := runMonitorCheckWithOptions(monitor, channelMonitorCheckOptions{
		client:      server.Client(),
		validateURL: allowTestChannelMonitorURL,
	})

	require.Error(t, err)
	var history model.ChannelMonitorHistory
	require.NoError(t, db.First(&history, "monitor_id = ?", monitor.Id).Error)
	require.Equal(t, "error", history.Status)
	require.Contains(t, history.ErrorMessage, "invalid custom_headers json object")
}