diff --git a/internal/core/dify_invocation/encrypt_test.go b/internal/core/dify_invocation/encrypt_test.go new file mode 100644 index 0000000..4eece3f --- /dev/null +++ b/internal/core/dify_invocation/encrypt_test.go @@ -0,0 +1,132 @@ +package dify_invocation + +import ( + "fmt" + "log" + "net/http" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/internal/utils/network" +) + +func TestEncryptRequired(t *testing.T) { + data := map[string]any{ + "key": "value", + } + + payload := &InvokeEncryptRequest{ + BaseInvokeDifyRequest: BaseInvokeDifyRequest{ + TenantId: "123", + UserId: "456", + Type: INVOKE_TYPE_ENCRYPT, + }, + InvokeEncryptSchema: InvokeEncryptSchema{ + Opt: ENCRYPT_OPT_ENCRYPT, + Namespace: ENCRYPT_NAMESPACE_ENDPOINT, + Identity: "test123", + Data: data, + Config: map[string]plugin_entities.ProviderConfig{ + "key": { + Name: "key", + Type: plugin_entities.CONFIG_TYPE_SECRET_INPUT, + }, + }, + }, + } + + if !payload.EncryptRequired(data) { + t.Errorf("EncryptRequired should return true") + } + + payload.Config["key"] = plugin_entities.ProviderConfig{ + Name: "key", + Type: plugin_entities.CONFIG_TYPE_TEXT_INPUT, + } + + if payload.EncryptRequired(data) { + t.Errorf("EncryptRequired should return false") + } +} + +func TestInvokeEncrypt(t *testing.T) { + server := gin.Default() + gin.SetMode(gin.ReleaseMode) + + port, err := network.GetRandomPort() + if err != nil { + t.Errorf("GetRandomPort failed: %v", err) + } + + http_invoked := false + + server.POST("/inner/api/invoke/encrypt", func(ctx *gin.Context) { + data := make(map[string]any) + if err := ctx.BindJSON(&data); err != nil { + t.Errorf("BindJSON failed: %v", err) + } + + if data["data"].(map[string]any)["key"] != "value" { + t.Errorf("data[key] should be `value`, but got %v", data["data"].(map[string]any)["key"]) + } + + http_invoked = true + + ctx.JSON(http.StatusOK, gin.H{ + "key": "encrypted", + }) + }) + + srv := &http.Server{ + Addr: fmt.Sprintf(":%d", port), + Handler: server, + } + + go func() { + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("listen: %s\n", err) + } + }() + + defer srv.Close() + + time.Sleep(1 * time.Second) + + if err := InitDifyInvocationDaemon(fmt.Sprintf("http://localhost:%d", port), "test"); err != nil { + t.Errorf("InitDifyInvocationDaemon failed: %v", err) + } + + payload := &InvokeEncryptRequest{ + BaseInvokeDifyRequest: BaseInvokeDifyRequest{ + TenantId: "123", + UserId: "456", + Type: INVOKE_TYPE_ENCRYPT, + }, + InvokeEncryptSchema: InvokeEncryptSchema{ + Opt: ENCRYPT_OPT_ENCRYPT, + Namespace: ENCRYPT_NAMESPACE_ENDPOINT, + Identity: "test123", + Data: map[string]any{"key": "value"}, + Config: map[string]plugin_entities.ProviderConfig{ + "key": { + Name: "key", + Type: plugin_entities.CONFIG_TYPE_SECRET_INPUT, + }, + }, + }, + } + + if encrypted, err := InvokeEncrypt(payload); err != nil { + t.Errorf("InvokeEncrypt failed: %v", err) + } else { + if encrypted["key"] != "encrypted" { + t.Errorf("encrypted[key] should be `encrypted`, but got %v", encrypted["key"]) + } + } + + if !http_invoked { + t.Errorf("http_invoked should be true") + } +}