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