From 696992f42dc1de196b52c81ecdaa93231a34e4e0 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Tue, 16 May 2023 13:59:05 -0700
Subject: [PATCH] Add SetDeadline test

---
 p2p/test/transport/deadline_test.go | 34 +++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/p2p/test/transport/deadline_test.go b/p2p/test/transport/deadline_test.go
index 3f79f1af70..64565fdadc 100644
--- a/p2p/test/transport/deadline_test.go
+++ b/p2p/test/transport/deadline_test.go
@@ -68,6 +68,11 @@ func TestReadWriteDeadlines(t *testing.T) {
 				require.Error(t, err)
 				require.Contains(t, err.Error(), "deadline")
 				require.Less(t, time.Since(start), 100*time.Millisecond)
+
+				if strings.Contains(tc.Name, "mplex") {
+					// FIXME: mplex stalls on close, so we reset so we don't spend an extra 5s waiting for nothing
+					s.Reset()
+				}
 			})
 
 			t.Run("WriteDeadline", func(t *testing.T) {
@@ -88,6 +93,35 @@ func TestReadWriteDeadlines(t *testing.T) {
 					s.Reset()
 				}
 			})
+
+			t.Run("Deadline for", func(t *testing.T) {
+				for _, op := range []string{"Read", "Write"} {
+					t.Run(op, func(t *testing.T) {
+						s, err := dialer.NewStream(context.Background(), listener.ID(), "/stall/1")
+						require.NoError(t, err)
+						defer s.Close()
+
+						// Set a deadline
+						s.SetDeadline(time.Now().Add(10 * time.Millisecond))
+						start := time.Now()
+
+						if op == "Read" {
+							buf := make([]byte, 1)
+							_, err = s.Read(buf)
+						} else {
+							_, err = s.Write(sendBuf)
+						}
+						require.Error(t, err)
+						require.Contains(t, err.Error(), "deadline")
+						require.Less(t, time.Since(start), 100*time.Millisecond)
+
+						if strings.Contains(tc.Name, "mplex") {
+							// FIXME: mplex stalls on close, so we reset so we don't spend an extra 5s waiting for nothing
+							s.Reset()
+						}
+					})
+				}
+			})
 		})
 	}
 }