From e24c08f444bb6da89283ab87e7aa0cc6cdc805e3 Mon Sep 17 00:00:00 2001
From: Dennis Trautwein <git@dtrautwein.eu>
Date: Sat, 1 May 2021 16:52:33 +0200
Subject: [PATCH] Defer closing stream for reading

---
 suites/transport/stream_suite.go    |  5 ++++-
 suites/transport/transport_suite.go | 20 ++++++++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/suites/transport/stream_suite.go b/suites/transport/stream_suite.go
index f7172b4..b1db0bc 100644
--- a/suites/transport/stream_suite.go
+++ b/suites/transport/stream_suite.go
@@ -55,7 +55,7 @@ type Options struct {
 }
 
 func fullClose(t *testing.T, s mux.MuxedStream) {
-	if err := s.Close(); err != nil {
+	if err := s.CloseWrite(); err != nil {
 		t.Error(err)
 		s.Reset()
 		return
@@ -67,6 +67,9 @@ func fullClose(t *testing.T, s mux.MuxedStream) {
 	if len(b) != 0 {
 		t.Error("expected to be done reading")
 	}
+	if err := s.Close(); err != nil {
+		t.Error(err)
+	}
 }
 
 func randBuf(size int) []byte {
diff --git a/suites/transport/transport_suite.go b/suites/transport/transport_suite.go
index 234c43c..c6a8c1c 100644
--- a/suites/transport/transport_suite.go
+++ b/suites/transport/transport_suite.go
@@ -141,7 +141,7 @@ func SubtestBasic(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr,
 		t.Fatalf("failed to write enough data (a->b)")
 		return
 	}
-	err = s.Close()
+	err = s.CloseWrite()
 	if err != nil {
 		t.Fatal(err)
 		return
@@ -155,6 +155,12 @@ func SubtestBasic(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr,
 	if !bytes.Equal(testData, buf) {
 		t.Errorf("expected %s, got %s", testData, buf)
 	}
+
+	err = s.Close()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
 }
 
 func SubtestPingPong(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) {
@@ -265,7 +271,11 @@ func SubtestPingPong(t *testing.T, ta, tb transport.Transport, maddr ma.Multiadd
 				t.Error("failed to write enough data (a->b)")
 				return
 			}
-			s.Close()
+			err = s.CloseWrite()
+			if err != nil {
+				t.Fatal(err)
+				return
+			}
 
 			ret, err := ioutil.ReadAll(s)
 			if err != nil {
@@ -276,6 +286,12 @@ func SubtestPingPong(t *testing.T, ta, tb transport.Transport, maddr ma.Multiadd
 			if !bytes.Equal(data, ret) {
 				t.Errorf("expected %q, got %q", string(data), string(ret))
 			}
+
+			err = s.Close()
+			if err != nil {
+				t.Fatal(err)
+				return
+			}
 		}(i)
 	}
 	wg.Wait()