|
@@ -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")
|
|
|
|
+ }
|
|
|
|
+}
|