refactor tests

Signed-off-by: Lixia (Sylvia) Lei <[email protected]>

internal/trace/transport.go

@@ -99,8 +99,6 @@ func logHeader(header http.Header) string {
 
 // logResponseBody prints out the response body if it is printable and within
 // the size limit.
-// TODO: review tests
-// TODO: what about seek case?
 func logResponseBody(resp *http.Response) string {
 	if resp.Body == nil || resp.Body == http.NoBody || resp.ContentLength == 0 {
 		return "   No response body to print"
@@ -118,12 +116,12 @@ func logResponseBody(resp *http.Response) string {
 		return fmt.Sprintf("   Response body larger than %d bytes is not printed", payloadSizeLimit)
 	}
 
-	// Note: Even if the actual body size mismatches the content length, we still print the body for debugging purposes.
-	// If the actual body size exceeds the limit, the body will be truncated to limit.
+	// Note: If the actual body size mismatches the content length and exceeds the limit,
+	// the body will be truncated to the limit for seucrity consideration.
 	// In this case, the response processing subsequent to logging might be broken.
-	// TODO: can we tell the body is truncated?
-	defer resp.Body.Close()
-	lr := io.LimitReader(resp.Body, payloadSizeLimit)
+	rc := resp.Body
+	defer rc.Close()
+	lr := io.LimitReader(rc, payloadSizeLimit)
 	bodyBytes, err := io.ReadAll(lr)
 	if err != nil {
 		return fmt.Sprintf("   Error reading response body: %v", err)

internal/trace/transport_test.go

@@ -16,10 +16,10 @@ limitations under the License.
 package trace
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"net/http"
-	"strings"
 	"testing"
 )
 
@@ -113,15 +113,15 @@ func Test_logResponseBody(t *testing.T) {
 			name: "No body",
 			resp: &http.Response{
 				Body:          http.NoBody,
-				ContentLength: 100,
+				ContentLength: 100, // in case of HEAD response, the content length is set but the body is empty
 				Header:        http.Header{"Content-Type": []string{"application/json"}},
 			},
 			want: "   No response body to print",
 		},
 		{
 			name: "Empty body",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("")),
+				Body:          io.NopCloser(bytes.NewReader([]byte(""))),
 				ContentLength: 0,
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
@@ -130,7 +130,7 @@ func Test_logResponseBody(t *testing.T) {
 		{
 			name: "Unknown content length",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("whatever")),
+				Body:          io.NopCloser(bytes.NewReader([]byte("whatever"))),
 				ContentLength: -1,
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
@@ -139,16 +139,25 @@ func Test_logResponseBody(t *testing.T) {
 		{
 			name: "Non-printable content type",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("binary data")),
-				ContentLength: 10,
+				Body:          io.NopCloser(bytes.NewReader([]byte("binary data"))),
+				ContentLength: 11,
 				Header:        http.Header{"Content-Type": []string{"application/octet-stream"}},
 			},
 			want: "   Response body of content type \"application/octet-stream\" is not printed",
 		},
+		{
+			name: "Body at the limit",
+			resp: &http.Response{
+				Body:          io.NopCloser(bytes.NewReader(bytes.Repeat([]byte("a"), int(payloadSizeLimit)))), // 1 byte larger than limit
+				ContentLength: payloadSizeLimit,
+				Header:        http.Header{"Content-Type": []string{"text/plain"}},
+			},
+			want: string(bytes.Repeat([]byte("a"), int(payloadSizeLimit))),
+		},
 		{
 			name: "Body larger than limit",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader(strings.Repeat("a", int(payloadSizeLimit+1)))), // 1 byte larger than limit
+				Body:          io.NopCloser(bytes.NewReader(bytes.Repeat([]byte("a"), int(payloadSizeLimit+1)))), // 1 byte larger than limit
 				ContentLength: payloadSizeLimit + 1,
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
@@ -157,16 +166,16 @@ func Test_logResponseBody(t *testing.T) {
 		{
 			name: "Printable content type within limit",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("data")),
-				ContentLength: int64(len("data")),
+				Body:          io.NopCloser(bytes.NewReader([]byte("data"))),
+				ContentLength: 4,
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
 			want: "data",
 		},
 		{
 			name: "Actual body size is larger than content length",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("data")),
+				Body:          io.NopCloser(bytes.NewReader([]byte("data"))),
 				ContentLength: 3, // mismatched content length
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
@@ -175,16 +184,16 @@ func Test_logResponseBody(t *testing.T) {
 		{
 			name: "Actual body size is larger than content length and exceeds limit",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader(strings.Repeat("a", int(payloadSizeLimit+1)))), // 1 byte larger than limit
-				ContentLength: 1,                                                                             // mismatched content length
+				Body:          io.NopCloser(bytes.NewReader(bytes.Repeat([]byte("a"), int(payloadSizeLimit+1)))), // 1 byte larger than limit
+				ContentLength: 1,                                                                                 // mismatched content length
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},
-			want: strings.Repeat("a", int(payloadSizeLimit)),
+			want: string(bytes.Repeat([]byte("a"), int(payloadSizeLimit))),
 		},
 		{
 			name: "Actual body size is smaller than content length",
 			resp: &http.Response{
-				Body:          io.NopCloser(strings.NewReader("data")),
+				Body:          io.NopCloser(bytes.NewReader([]byte("data"))),
 				ContentLength: 5, // mismatched content length
 				Header:        http.Header{"Content-Type": []string{"text/plain"}},
 			},