martian(closeWriter): check closeWriter recursively

saucelabs · Mar 4, 2024 · e3a6034 · e3a6034
1 parent 6aebf00
commit e3a6034
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/internal/martian/closewriter.go b/internal/martian/closewriter.go
@@ -44,6 +44,8 @@ func asCloseWriter(w io.Writer) (closeWriter, bool) {
 	if v.Kind() != reflect.Struct {
 		return nil, false
 	}
+
+	// Check if any of the fields implement closeWriter.
 	for i := 0; i < v.NumField(); i++ {
 		f := v.Field(i)
 		if f.CanInterface() {
@@ -53,5 +55,17 @@ func asCloseWriter(w io.Writer) (closeWriter, bool) {
 		}
 	}
 
+	// Check recursively...
+	for i := 0; i < v.NumField(); i++ {
+		f := v.Field(i)
+		if f.CanInterface() {
+			if w, ok := f.Interface().(io.Writer); ok {
+				if cw, ok := asCloseWriter(w); ok {
+					return cw, true
+				}
+			}
+		}
+	}
+
 	return nil, false
 }
diff --git a/internal/martian/closewriter_test.go b/internal/martian/closewriter_test.go
@@ -106,6 +106,30 @@ func TestAsCloseWriter(t *testing.T) {
 				&nopCloseWriterPtrImpl{},
 			},
 		},
+		{
+			name: "embedded nopCloseWriterImpl",
+			w: struct {
+				io.Writer
+			}{
+				struct {
+					io.Writer
+				}{
+					nopCloseWriterImpl{},
+				},
+			},
+		},
+		{
+			name: "embedded ptr nopCloseWriterImpl",
+			w: struct {
+				io.Writer
+			}{
+				&struct {
+					io.Writer
+				}{
+					nopCloseWriterImpl{},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {