aiops/backend/middleware/permission_test.go

package middleware

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

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
)

func TestPermissionAuth(t *testing.T) {
	gin.SetMode(gin.TestMode)
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	if err != nil {
		t.Fatalf("打开测试数据库失败: %v", err)
	}
	statements := []string{
		"CREATE TABLE iam_user_roles (user_id INTEGER, role_id INTEGER)",
		"CREATE TABLE iam_roles (id INTEGER PRIMARY KEY, status INTEGER)",
		"CREATE TABLE iam_role_permissions (role_id INTEGER, permission_id INTEGER)",
		"CREATE TABLE iam_permissions (id INTEGER PRIMARY KEY, code TEXT, status INTEGER)",
		"INSERT INTO iam_user_roles(user_id, role_id) VALUES (7, 1)",
		"INSERT INTO iam_roles(id, status) VALUES (1, 1)",
		"INSERT INTO iam_role_permissions(role_id, permission_id) VALUES (1, 10)",
		"INSERT INTO iam_permissions(id, code, status) VALUES (10, 'aiops:agent-cli', 1)",
	}
	for _, statement := range statements {
		if err := db.Exec(statement).Error; err != nil {
			t.Fatalf("准备测试数据失败: %v", err)
		}
	}

	tests := []struct {
		name       string
		userID     uint
		isAdmin    bool
		wantStatus int
	}{
		{name: "granted", userID: 7, wantStatus: http.StatusNoContent},
		{name: "denied", userID: 8, wantStatus: http.StatusForbidden},
		{name: "admin", userID: 9, isAdmin: true, wantStatus: http.StatusNoContent},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			router := gin.New()
			router.Use(func(c *gin.Context) {
				c.Set("userID", tc.userID)
				c.Set("isAdmin", tc.isAdmin)
			})
			router.GET("/agent", PermissionAuth(db, "aiops:agent-cli"), func(c *gin.Context) {
				c.Status(http.StatusNoContent)
			})

			recorder := httptest.NewRecorder()
			router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/agent", nil))
			if recorder.Code != tc.wantStatus {
				t.Fatalf("状态码=%d, want %d, body=%s", recorder.Code, tc.wantStatus, recorder.Body.String())
			}
		})
	}
}