new-api/controller/muse_internal_test.go

package controller

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

	"github.com/QuantumNous/new-api/common"
	museDTO "github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

type museInternalAPIResponse struct {
	Success bool                         `json:"success"`
	Message string                       `json:"message"`
	Data    museDTO.SyncMuseUserResponse `json:"data"`
}

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

	gin.SetMode(gin.TestMode)
	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	common.RedisEnabled = false
	common.GlobalApiRateLimitEnable = false

	t.Setenv("MUSE_INTERNAL_SECRET", "test-secret")

	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)

	model.DB = db
	model.LOG_DB = db
	require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}))

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

func decodeMuseInternalResponse(t *testing.T, recorder *httptest.ResponseRecorder) museInternalAPIResponse {
	t.Helper()

	var response museInternalAPIResponse
	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
	return response
}

func newMuseInternalRouter() *gin.Engine {
	engine := gin.New()
	internal := engine.Group("/api/internal/muse")
	internal.Use(middleware.MuseInternalAuth())
	{
		internal.POST("/users/sync", SyncMuseUser)
		internal.POST("/users/:id/disable", DisableMuseUser)
		internal.POST("/tokens/:id/revoke", RevokeMuseToken)
		internal.GET("/users/:id/status", GetMuseUserStatus)
	}
	return engine
}

func TestMuseInternalSyncUserCreatesMirrorUserAndToken(t *testing.T) {
	setupMuseInternalControllerTestDB(t)

	req := httptest.NewRequest(http.MethodPost, "/api/internal/muse/users/sync", strings.NewReader(`{"muse_user_id":"u-1","username":"alice","display_name":"Alice","status":"active"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Muse-Service-Secret", "test-secret")
	rr := httptest.NewRecorder()

	newMuseInternalRouter().ServeHTTP(rr, req)

	require.Equal(t, http.StatusOK, rr.Code)
	response := decodeMuseInternalResponse(t, rr)
	require.True(t, response.Success)
	require.Equal(t, "active", response.Data.Status)
	require.NotZero(t, response.Data.NewAPIUserID)
	require.NotZero(t, response.Data.TokenID)
	require.NotEmpty(t, response.Data.TokenKey)

	user, err := model.GetUserById(response.Data.NewAPIUserID, false)
	require.NoError(t, err)
	require.Equal(t, common.UserStatusEnabled, user.Status)
	require.Equal(t, "alice", user.Username)
	require.Equal(t, "muse_user_id:u-1", user.Remark)

	token, err := model.GetTokenById(response.Data.TokenID)
	require.NoError(t, err)
	require.Equal(t, user.Id, token.UserId)
	require.Equal(t, common.TokenStatusEnabled, token.Status)
	require.True(t, token.UnlimitedQuota)
}

func TestMuseInternalSyncUserRejectsInvalidSecret(t *testing.T) {
	setupMuseInternalControllerTestDB(t)

	req := httptest.NewRequest(http.MethodPost, "/api/internal/muse/users/sync", strings.NewReader(`{"muse_user_id":"u-1","username":"alice","display_name":"Alice","status":"active"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Muse-Service-Secret", "wrong-secret")
	rr := httptest.NewRecorder()

	newMuseInternalRouter().ServeHTTP(rr, req)

	require.Equal(t, http.StatusUnauthorized, rr.Code)
}

func TestMuseInternalDisableUserDisablesUserAndToken(t *testing.T) {
	setupMuseInternalControllerTestDB(t)

	syncReq := httptest.NewRequest(http.MethodPost, "/api/internal/muse/users/sync", strings.NewReader(`{"muse_user_id":"u-2","username":"bob","display_name":"Bob","status":"active"}`))
	syncReq.Header.Set("Content-Type", "application/json")
	syncReq.Header.Set("X-Muse-Service-Secret", "test-secret")
	syncRR := httptest.NewRecorder()
	engine := newMuseInternalRouter()
	engine.ServeHTTP(syncRR, syncReq)

	syncResponse := decodeMuseInternalResponse(t, syncRR)
	require.True(t, syncResponse.Success)

	disableReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/internal/muse/users/%d/disable", syncResponse.Data.NewAPIUserID), nil)
	disableReq.Header.Set("X-Muse-Service-Secret", "test-secret")
	disableRR := httptest.NewRecorder()
	engine.ServeHTTP(disableRR, disableReq)

	require.Equal(t, http.StatusOK, disableRR.Code)
	disableResponse := decodeMuseInternalResponse(t, disableRR)
	require.True(t, disableResponse.Success)
	require.Equal(t, "disabled", disableResponse.Data.Status)
	require.Equal(t, syncResponse.Data.NewAPIUserID, disableResponse.Data.NewAPIUserID)

	user, err := model.GetUserById(syncResponse.Data.NewAPIUserID, false)
	require.NoError(t, err)
	require.Equal(t, common.UserStatusDisabled, user.Status)

	token, err := model.GetTokenById(syncResponse.Data.TokenID)
	require.NoError(t, err)
	require.Equal(t, common.TokenStatusDisabled, token.Status)
}