Pārlūkot izejas kodu

feat: add test for redirecting traffic with query parameters and refactor URL construction logic

Yeuoly 6 mēneši atpakaļ
vecāks
revīzija
9e676b7468
2 mainītis faili ar 29 papildinājumiem un 4 dzēšanām
  1. 9 4
      internal/cluster/redirect.go
  2. 20 0
      internal/cluster/redirect_test.go

+ 9 - 4
internal/cluster/redirect.go

@@ -6,6 +6,14 @@ import (
 	"net/http"
 )
 
+func constructRedirectUrl(ip address, request *http.Request) string {
+	url := "http://" + ip.fullAddress() + request.URL.Path
+	if request.URL.RawQuery != "" {
+		url += "?" + request.URL.RawQuery
+	}
+	return url
+}
+
 // RedirectRequest redirects the request to the specified node
 func (c *Cluster) RedirectRequest(
 	node_id string, request *http.Request,
@@ -22,10 +30,7 @@ func (c *Cluster) RedirectRequest(
 
 	ip := ips[0]
 
-	url := "http://" + ip.fullAddress() + request.URL.Path
-	if request.URL.RawQuery != "" {
-		url += "?" + request.URL.RawQuery
-	}
+	url := constructRedirectUrl(ip, request)
 
 	// create a new request
 	redirectedRequest, err := http.NewRequest(

+ 20 - 0
internal/cluster/redirect_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/network"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/endpoint_entities"
 )
 
 type SimulationCheckServer struct {
@@ -180,3 +181,22 @@ func TestRedirectTraffic(t *testing.T) {
 		}
 	}
 }
+
+func TestRedirectTrafficWithQueryParams(t *testing.T) {
+	request, err := http.NewRequest("GET", "http://localhost:8080/plugin/invoke/tool?a=1&b=2", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	request.Header.Set(endpoint_entities.HeaderXOriginalHost, "localhost:8080")
+
+	ip := address{
+		Ip:   "127.0.0.1",
+		Port: 8080,
+	}
+
+	redirectedRequest := constructRedirectUrl(ip, request)
+	if redirectedRequest != "http://127.0.0.1:8080/plugin/invoke/tool?a=1&b=2" {
+		t.Fatal("redirected request is not correct")
+	}
+}