From 39006fcf1623eb8a946dcd1b2ed92d9c5148fb52 Mon Sep 17 00:00:00 2001
From: TimAndy <xuchonglei@126.com>
Date: Thu, 3 Aug 2023 10:28:56 +0800
Subject: [PATCH] test: Add custom transport test case (#1601)

---
 client_test.go | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/client_test.go b/client_test.go
index 6da830416d..64e709a14d 100644
--- a/client_test.go
+++ b/client_test.go
@@ -3192,3 +3192,61 @@ func Test_AddMissingPort(t *testing.T) {
 		})
 	}
 }
+
+type TransportEx struct {
+	base    RoundTripper
+	reqLog  *string
+	respLog *string
+}
+
+func (t TransportEx) RoundTrip(hc *HostClient, req *Request, resp *Response) (bool, error) {
+	req.Header.Set("trace-id", "123")
+	*t.reqLog = req.String()
+	retry, err := t.base.RoundTrip(hc, req, resp)
+	*t.respLog = resp.String()
+	return retry, err
+}
+
+func TestClientTransportEx(t *testing.T) {
+	var reqLog string
+	var respLog string
+
+	sHTTP := startEchoServer(t, "tcp", "127.0.0.1:")
+	defer sHTTP.Stop()
+
+	sHTTPS := startEchoServerTLS(t, "tcp", "127.0.0.1:")
+	defer sHTTPS.Stop()
+
+	c := &Client{
+		TLSConfig: &tls.Config{
+			InsecureSkipVerify: true,
+		},
+		ConfigureClient: func(hc *HostClient) error {
+			hc.Transport = TransportEx{base: DefaultTransport, reqLog: &reqLog, respLog: &respLog}
+			return nil
+		},
+	}
+
+	for i := 0; i < 4; i++ {
+		addr := "http://" + sHTTP.Addr()
+		if i&1 != 0 {
+			addr = "https://" + sHTTPS.Addr()
+		}
+		// test get
+		testClientGet(t, c, addr, 20)
+		if !strings.Contains(reqLog, "Trace-Id: 123") {
+			t.Errorf("request log should contains: %v", "Trace-Id: 123")
+		}
+		if !strings.Contains(respLog, "HTTP/1.1 200 OK") {
+			t.Errorf("response log should contains: %v", "HTTP/1.1 200 OK")
+		}
+		// test post
+		testClientPost(t, c, addr, 10)
+		if !strings.Contains(reqLog, "Trace-Id: 123") {
+			t.Errorf("request log should contains: %v", "Trace-Id: 123")
+		}
+		if !strings.Contains(respLog, "HTTP/1.1 200 OK") {
+			t.Errorf("response log should contains: %v", "HTTP/1.1 200 OK")
+		}
+	}
+}