瀏覽代碼

feat: add tests for permission checker

Yeuoly 1 年之前
父節點
當前提交
126322f75f

+ 10 - 10
internal/core/plugin_daemon/backwards_invocation/task.go

@@ -58,49 +58,49 @@ func InvokeDify(
 var (
 	permissionMapping = map[dify_invocation.InvokeType]map[string]any{
 		dify_invocation.INVOKE_TYPE_TOOL: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeTool()
 			},
 			"error": "permission denied, you need to enable tool access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_LLM: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeLLM()
 			},
 			"error": "permission denied, you need to enable llm access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_TEXT_EMBEDDING: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeTextEmbedding()
 			},
 			"error": "permission denied, you need to enable text-embedding access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_RERANK: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeRerank()
 			},
 			"error": "permission denied, you need to enable rerank access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_TTS: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeTTS()
 			},
 			"error": "permission denied, you need to enable tts access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_SPEECH2TEXT: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeSpeech2Text()
 			},
 			"error": "permission denied, you need to enable speech2text access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_MODERATION: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeModeration()
 			},
 			"error": "permission denied, you need to enable moderation access in plugin manifest",
 		},
 		dify_invocation.INVOKE_TYPE_NODE: {
-			"func": func(runtime entities.PluginRuntimeInterface) bool {
+			"func": func(runtime entities.PluginRuntimeTimeLifeInterface) bool {
 				return runtime.Configuration().Resource.Permission.AllowInvokeNode()
 			},
 			"error": "permission denied, you need to enable node access in plugin manifest",
@@ -108,13 +108,13 @@ var (
 	}
 )
 
-func checkPermission(runtime entities.PluginRuntimeInterface, request_handle *BackwardsInvocation) error {
+func checkPermission(runtime entities.PluginRuntimeTimeLifeInterface, request_handle *BackwardsInvocation) error {
 	permission, ok := permissionMapping[request_handle.Type()]
 	if !ok {
 		return fmt.Errorf("unsupported invoke type: %s", request_handle.Type())
 	}
 
-	permission_func, ok := permission["func"].(func(runtime entities.PluginRuntimeInterface) bool)
+	permission_func, ok := permission["func"].(func(runtime entities.PluginRuntimeTimeLifeInterface) bool)
 	if !ok {
 		return fmt.Errorf("permission function not found: %s", request_handle.Type())
 	}

+ 148 - 0
internal/core/plugin_daemon/backwards_invocation/task_test.go

@@ -0,0 +1,148 @@
+package backwards_invocation
+
+import (
+	"testing"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation"
+	"github.com/langgenius/dify-plugin-daemon/internal/types/entities"
+	"github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
+)
+
+type TPluginRuntime struct {
+	entities.PluginRuntime
+}
+
+func (r *TPluginRuntime) InitEnvironment() error {
+	return nil
+}
+
+func (r *TPluginRuntime) StartPlugin() error {
+	return nil
+}
+
+func (r *TPluginRuntime) Type() entities.PluginRuntimeType {
+	return entities.PLUGIN_RUNTIME_TYPE_LOCAL
+}
+
+func (r *TPluginRuntime) Wait() (<-chan bool, error) {
+	return nil, nil
+}
+
+func TestBackwardsInvocationAllPermittedPermission(t *testing.T) {
+	all_permitted_runtime := TPluginRuntime{
+		PluginRuntime: entities.PluginRuntime{
+			Config: plugin_entities.PluginDeclaration{
+				Resource: plugin_entities.PluginResourceRequirement{
+					Permission: &plugin_entities.PluginPermissionRequirement{
+						Tool: &plugin_entities.PluginPermissionToolRequirement{
+							Enabled: true,
+						},
+						Model: &plugin_entities.PluginPermissionModelRequirement{
+							Enabled:       true,
+							LLM:           true,
+							TextEmbedding: true,
+							Rerank:        true,
+							Moderation:    true,
+							TTS:           true,
+							Speech2text:   true,
+						},
+						Node: &plugin_entities.PluginPermissionNodeRequirement{
+							Enabled: true,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	invoke_llm_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_LLM, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_llm_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_text_embedding_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TEXT_EMBEDDING, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_text_embedding_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_rerank_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_RERANK, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_rerank_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_tts_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TTS, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_tts_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_speech2text_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_SPEECH2TEXT, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_speech2text_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_moderation_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_MODERATION, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_moderation_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_tool_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TOOL, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_tool_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+
+	invoke_node_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_NODE, "", nil, nil)
+	if err := checkPermission(&all_permitted_runtime, invoke_node_request); err != nil {
+		t.Errorf("checkPermission failed: %s", err.Error())
+	}
+}
+
+func TestBackwardsInvocationAllDeniedPermission(t *testing.T) {
+	all_denied_runtime := TPluginRuntime{
+		PluginRuntime: entities.PluginRuntime{
+			Config: plugin_entities.PluginDeclaration{
+				Resource: plugin_entities.PluginResourceRequirement{},
+			},
+		},
+	}
+
+	invoke_llm_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_LLM, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_llm_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_text_embedding_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TEXT_EMBEDDING, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_text_embedding_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_rerank_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_RERANK, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_rerank_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_tts_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TTS, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_tts_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_speech2text_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_SPEECH2TEXT, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_speech2text_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_moderation_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_MODERATION, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_moderation_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_tool_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_TOOL, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_tool_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+	invoke_node_request := NewBackwardsInvocation(dify_invocation.INVOKE_TYPE_NODE, "", nil, nil)
+	if err := checkPermission(&all_denied_runtime, invoke_node_request); err == nil {
+		t.Errorf("checkPermission failed: expected error, got nil")
+	}
+
+}