guard/internal/usecase/opa_test.go

package usecase

import (
	"ark-permission/internal/domain/usecase"
	"context"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zeromicro/go-zero/core/logx"
	"testing"
	"time"
)

func TestMustOpaUseCase(t *testing.T) {
	// 初始化 OPA UseCase
	got, err := NewOpaUseCase(OpaUseCaseParam{})
	assert.NoError(t, err)

	ctx := context.Background()

	// 加载 Policy
	err = got.LoadPolicy(ctx, []usecase.Policy{
		{
			Role:    "admin",
			Path:    "/admin/.*",
			Methods: []string{"GET", "POST"},
			Name:    "Admin access",
		},
		{
			Role:    "user",
			Path:    "/user/.*",
			Methods: []string{"GET"},
			Name:    "User read access",
		},
		{
			Role:    "editor",
			Path:    "/editor/.*",
			Methods: []string{"PUT", "POST"},
			Name:    "Editor access",
		},
	})
	assert.NoError(t, err)

	// 定義測試表
	tests := []struct {
		name        string
		req         usecase.CheckReq
		expect      bool
		expectError bool
	}{
		{
			name: "單一角色，應允許通過",
			req: usecase.CheckReq{
				Roles:  []string{"user"},
				Path:   "/user/profile",
				Method: "GET",
			},
			expect: true,
		},
		{
			name: "多角色其中一個有配到，應允許通過",
			req: usecase.CheckReq{
				Roles:  []string{"user", "admin"},
				Path:   "/user/profile",
				Method: "GET",
			},
			expect: true,
		},
		{
			name: "角色不匹配，應拒絕通過",
			req: usecase.CheckReq{
				Roles:  []string{"editor"},
				Path:   "/user/profile",
				Method: "GET",
			},
			expect: false,
		},
		{
			name: "路徑不匹配，應拒絕通過",
			req: usecase.CheckReq{
				Roles:  []string{"user"},
				Path:   "/editor/dashboard",
				Method: "GET",
			},
			expect: false,
		},
		{
			name: "方法不匹配，應拒絕通過",
			req: usecase.CheckReq{
				Roles:  []string{"user"},
				Path:   "/user/profile",
				Method: "POST",
			},
			expect: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			check, err := got.CheckRBACPermission(ctx, tt.req)
			if tt.expectError {
				assert.Error(t, err, "expected an error but got none")
			} else {
				assert.NoError(t, err, "did not expect an error but got one")
				assert.Equal(t, tt.expect, check.Allow)
			}
		})
	}
}

func TestLoadPolicy(t *testing.T) {
	// 初始化 OPA UseCase
	got, err := NewOpaUseCase(OpaUseCaseParam{})
	require.NoError(t, err)

	tests := []struct {
		name       string
		input      []usecase.Policy
		ctxTimeout time.Duration
		expectErr  bool
	}{
		{
			name: "正常加載多個Policy",
			input: []usecase.Policy{
				{
					Role:    "admin",
					Path:    "/admin/.*",
					Methods: []string{"GET", "POST"},
					Name:    "Admin access",
				},
				{
					Role:    "user",
					Path:    "/user/.*",
					Methods: []string{"GET"},
					Name:    "User read access",
				},
			},
			ctxTimeout: 3 * time.Second, // 足夠的時間來執行
			expectErr:  false,
		},
		{
			name: "加載策略超時",
			input: []usecase.Policy{
				{
					Role:    "admin",
					Path:    "/admin/.*",
					Methods: []string{"GET", "POST"},
					Name:    "Admin access",
				},
				{
					Role:    "user",
					Path:    "/user/.*",
					Methods: []string{"GET"},
					Name:    "User read access",
				},
			},
			ctxTimeout: 1 * time.Nanosecond, // 超時
			expectErr:  true,
		},
		{
			name:       "空策略加載",
			input:      []usecase.Policy{},
			ctxTimeout: 3 * time.Second, // 足夠的時間
			expectErr:  false,
		},
	}

	// 遍歷所有測試用例
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 設置具有超時的 Context
			ctx, cancel := context.WithTimeout(context.Background(), tt.ctxTimeout)
			defer cancel()

			// 調用 LoadPolicy
			err := got.LoadPolicy(ctx, tt.input)

			// 檢查是否符合預期錯誤
			if tt.expectErr {
				assert.Error(t, err, "預期發生錯誤，但沒有發生")
			} else {
				assert.NoError(t, err, "不預期發生錯誤，但卻發生了")
			}

			// 如果沒有錯誤，檢查 policies 是否被正確加載
			if !tt.expectErr {
				assert.Equal(t, len(tt.input), len(got.GetPolicy(ctx)), "policies 加載的數量與輸入數量不一致")
			}
		})
	}
}

func BenchmarkLoadPolicy(b *testing.B) {
	// 初始化 OPA UseCase
	got, _ := NewOpaUseCase(OpaUseCaseParam{})
	logx.Disable()

	policiesSmall := []usecase.Policy{
		{
			Role:    "admin",
			Path:    "/admin/.*",
			Methods: []string{"GET", "POST"},
			Name:    "Admin access",
		},
		{
			Role:    "user",
			Path:    "/user/.*",
			Methods: []string{"GET"},
			Name:    "User read access",
		},
	}

	policiesLarge := make([]usecase.Policy, 1000)
	for i := 0; i < 1000; i++ {
		policiesLarge[i] = usecase.Policy{
			Role:    "admin",
			Path:    "/admin/.*",
			Methods: []string{"GET", "POST"},
			Name:    "Admin access",
		}
	}

	benchmarks := []struct {
		name       string
		policies   []usecase.Policy
		ctxTimeout time.Duration
	}{
		{
			name:       "SmallPolicy_NoTimeout",
			policies:   policiesSmall,
			ctxTimeout: 1 * time.Second, // 沒有超時
		},
		{
			name:       "LargePolicy_NoTimeout",
			policies:   policiesLarge,
			ctxTimeout: 1 * time.Second, // 沒有超時
		},
		{
			name:       "SmallPolicy_WithTimeout",
			policies:   policiesSmall,
			ctxTimeout: 1 * time.Millisecond, // 沒有超時
		},
		{
			name:       "LargePolicy_WithTimeout",
			policies:   policiesLarge,
			ctxTimeout: 1 * time.Millisecond, // 沒有超時
		},
	}

	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				ctx, cancel := context.WithTimeout(context.Background(), bm.ctxTimeout)
				defer cancel()

				_ = got.LoadPolicy(ctx, bm.policies)
			}
		})
	}
}

func BenchmarkCheckRBACPermission(b *testing.B) {
	got, _ := NewOpaUseCase(OpaUseCaseParam{})
	logx.Disable()
	// 定義測試用 Policy
	policies := []usecase.Policy{
		{
			Role:    "admin",
			Path:    "/admin/.*",
			Methods: []string{"GET", "POST"},
			Name:    "Admin access",
		},
		{
			Role:    "user",
			Path:    "/user/.*",
			Methods: []string{"GET"},
			Name:    "User read access",
		},
		{
			Role:    "editor",
			Path:    "/editor/.*",
			Methods: []string{"PUT", "POST"},
			Name:    "Editor access",
		},
	}

	// 加載 Policy
	_ = got.LoadPolicy(context.Background(), policies)

	// 定義不同測試基準場景
	benchmarks := []struct {
		name string
		req  usecase.CheckReq
	}{
		{
			name: "SingleRole_SimplePath",
			req: usecase.CheckReq{
				Roles:  []string{"user"},
				Path:   "/user/profile",
				Method: "GET",
			},
		},
		{
			name: "MultipleRoles_ComplexPath",
			req: usecase.CheckReq{
				Roles:  []string{"admin", "user", "editor"},
				Path:   "/editor/dashboard",
				Method: "PUT",
			},
		},
		{
			name: "NoRoles_InvalidPath",
			req: usecase.CheckReq{
				Roles:  []string{},
				Path:   "/invalid/path",
				Method: "POST",
			},
		},
	}

	// 走訪所有場景
	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				// 設置一個超時的 ctx
				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
				defer cancel()

				_, _ = got.CheckRBACPermission(ctx, bm.req)
			}
		})
	}
}