capture and split notification if it has multiple hosts

pkg/proxy/remotehost.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -39,6 +40,33 @@ var (
 	logger = otelslog.NewLogger(name)
 )
 
+// MockResponseWriter implements http.ResponseWriter to capture responses for proxy calls.
+type MockResponseWriter struct {
+	StatusCode int
+	HeaderMap  http.Header
+	Body       *bytes.Buffer
+}
+
+func NewMockResponseWriter() *MockResponseWriter {
+	return &MockResponseWriter{
+		StatusCode: 200,
+		HeaderMap:  make(http.Header),
+		Body:       &bytes.Buffer{},
+	}
+}
+
+func (m *MockResponseWriter) Header() http.Header {
+	return m.HeaderMap
+}
+
+func (m *MockResponseWriter) Write(data []byte) (int, error) {
+	return m.Body.Write(data)
+}
+
+func (m *MockResponseWriter) WriteHeader(statusCode int) {
+	m.StatusCode = statusCode
+}
+
 func NewRemoteHost(host string) (*RemoteHost, error) {
 
 	target := fmt.Sprintf("%s:1337", host)
@@ -182,7 +210,7 @@ func (h *RemoteHost) AnnounceExpiry(uids []uint64) {
 	h.missedPings = 0
 }
 
-func (h *RemoteHost) Proxy(id uint64, w http.ResponseWriter, r *http.Request) (bool, error) {
+func (h *RemoteHost) Proxy(id uint64, w http.ResponseWriter, r *http.Request, customBody io.Reader) (bool, error) {
 	target := fmt.Sprintf("%s%s", h.httpBase, r.URL.RequestURI())
 
 	ctx, span := tracer.Start(r.Context(), "remote_proxy")
@@ -196,7 +224,11 @@ func (h *RemoteHost) Proxy(id uint64, w http.ResponseWriter, r *http.Request) (b
 	)
 	logger.InfoContext(ctx, "proxying request", "cartid", id, "host", h.host, "method", r.Method)
 
-	req, err := http.NewRequestWithContext(ctx, r.Method, target, r.Body)
+	var bdy io.Reader = r.Body
+	if customBody != nil {
+		bdy = customBody
+	}
+	req, err := http.NewRequestWithContext(ctx, r.Method, target, bdy)
 	if err != nil {
 		span.RecordError(err)
 		http.Error(w, "proxy build error", http.StatusBadGateway)