encrypt_test.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. package real
  2. import (
  3. "fmt"
  4. "log"
  5. "net/http"
  6. "testing"
  7. "time"
  8. "github.com/gin-gonic/gin"
  9. "github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation"
  10. "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
  11. "github.com/langgenius/dify-plugin-daemon/internal/utils/network"
  12. )
  13. func TestEncryptRequired(t *testing.T) {
  14. data := map[string]any{
  15. "key": "value",
  16. }
  17. payload := &dify_invocation.InvokeEncryptRequest{
  18. BaseInvokeDifyRequest: dify_invocation.BaseInvokeDifyRequest{
  19. TenantId: "123",
  20. UserId: "456",
  21. Type: dify_invocation.INVOKE_TYPE_ENCRYPT,
  22. },
  23. InvokeEncryptSchema: dify_invocation.InvokeEncryptSchema{
  24. Opt: dify_invocation.ENCRYPT_OPT_ENCRYPT,
  25. Namespace: dify_invocation.ENCRYPT_NAMESPACE_ENDPOINT,
  26. Identity: "test123",
  27. Data: data,
  28. Config: map[string]plugin_entities.ProviderConfig{
  29. "key": {
  30. Name: "key",
  31. Type: plugin_entities.CONFIG_TYPE_SECRET_INPUT,
  32. },
  33. },
  34. },
  35. }
  36. if !payload.EncryptRequired(data) {
  37. t.Errorf("EncryptRequired should return true")
  38. }
  39. payload.Config["key"] = plugin_entities.ProviderConfig{
  40. Name: "key",
  41. Type: plugin_entities.CONFIG_TYPE_TEXT_INPUT,
  42. }
  43. if payload.EncryptRequired(data) {
  44. t.Errorf("EncryptRequired should return false")
  45. }
  46. }
  47. func TestInvokeEncrypt(t *testing.T) {
  48. server := gin.Default()
  49. gin.SetMode(gin.ReleaseMode)
  50. port, err := network.GetRandomPort()
  51. if err != nil {
  52. t.Errorf("GetRandomPort failed: %v", err)
  53. }
  54. http_invoked := false
  55. server.POST("/inner/api/invoke/encrypt", func(ctx *gin.Context) {
  56. data := make(map[string]any)
  57. if err := ctx.BindJSON(&data); err != nil {
  58. t.Errorf("BindJSON failed: %v", err)
  59. }
  60. if data["data"].(map[string]any)["key"] != "value" {
  61. t.Errorf("data[key] should be `value`, but got %v", data["data"].(map[string]any)["key"])
  62. }
  63. http_invoked = true
  64. ctx.JSON(http.StatusOK, gin.H{
  65. "key": "encrypted",
  66. })
  67. })
  68. srv := &http.Server{
  69. Addr: fmt.Sprintf(":%d", port),
  70. Handler: server,
  71. }
  72. go func() {
  73. if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
  74. log.Fatalf("listen: %s\n", err)
  75. }
  76. }()
  77. defer srv.Close()
  78. time.Sleep(1 * time.Second)
  79. i, err := InitDifyInvocationDaemon(fmt.Sprintf("http://localhost:%d", port), "test")
  80. if err != nil {
  81. t.Errorf("InitDifyInvocationDaemon failed: %v", err)
  82. return
  83. }
  84. payload := &dify_invocation.InvokeEncryptRequest{
  85. BaseInvokeDifyRequest: dify_invocation.BaseInvokeDifyRequest{
  86. TenantId: "123",
  87. UserId: "456",
  88. Type: dify_invocation.INVOKE_TYPE_ENCRYPT,
  89. },
  90. InvokeEncryptSchema: dify_invocation.InvokeEncryptSchema{
  91. Opt: dify_invocation.ENCRYPT_OPT_ENCRYPT,
  92. Namespace: dify_invocation.ENCRYPT_NAMESPACE_ENDPOINT,
  93. Identity: "test123",
  94. Data: map[string]any{"key": "value"},
  95. Config: map[string]plugin_entities.ProviderConfig{
  96. "key": {
  97. Name: "key",
  98. Type: plugin_entities.CONFIG_TYPE_SECRET_INPUT,
  99. },
  100. },
  101. },
  102. }
  103. if encrypted, err := i.InvokeEncrypt(payload); err != nil {
  104. t.Errorf("InvokeEncrypt failed: %v", err)
  105. } else {
  106. if encrypted["key"] != "encrypted" {
  107. t.Errorf("encrypted[key] should be `encrypted`, but got %v", encrypted["key"])
  108. }
  109. }
  110. if !http_invoked {
  111. t.Errorf("http_invoked should be true")
  112. }
  113. }