From f90b62c3ce74143fb49486e32777ba418d1f05a5 Mon Sep 17 00:00:00 2001
From: dottyjones <38705838+dottyjones@users.noreply.github.com>
Date: Mon, 28 May 2018 07:28:04 -0700
Subject: [PATCH] Add test for handshake deadline

---
 client_server_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/client_server_test.go b/client_server_test.go
index dcaab3d..20b648b 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -344,6 +344,63 @@ func TestDialTimeout(t *testing.T) {
 	}
 }
 
+// netConnDeadlineObserver fails test if read or write called without deadline.
+type netConnDeadlineObserver struct {
+	t           *testing.T
+	c           net.Conn
+	read, write bool
+}
+
+func (c *netConnDeadlineObserver) SetDeadline(t time.Time) error {
+	c.write = !t.Equal(time.Time{})
+	c.read = c.write
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) SetReadDeadline(t time.Time) error {
+	c.read = !t.Equal(time.Time{})
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) SetWriteDeadline(t time.Time) error {
+	c.write = !t.Equal(time.Time{})
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) Write(p []byte) (int, error) {
+	if !c.write {
+		c.t.Fatalf("write with no deadline")
+	}
+	return c.c.Write(p)
+}
+
+func (c *netConnDeadlineObserver) Read(p []byte) (int, error) {
+	if !c.read {
+		c.t.Fatalf("read with no deadline")
+	}
+	return c.c.Read(p)
+}
+
+func (c *netConnDeadlineObserver) Close() error         { return c.c.Close() }
+func (c *netConnDeadlineObserver) LocalAddr() net.Addr  { return c.c.LocalAddr() }
+func (c *netConnDeadlineObserver) RemoteAddr() net.Addr { return c.c.RemoteAddr() }
+
+func TestHandshakeTimeout(t *testing.T) {
+	s := newServer(t)
+	defer s.Close()
+
+	d := cstDialer
+	d.NetDial = func(n, a string) (net.Conn, error) {
+		c, err := net.Dial(n, a)
+		return &netConnDeadlineObserver{c: c, t: t}, err
+	}
+	ws, _, err := d.Dial(s.URL, nil)
+	if err != nil {
+		t.Fatal("Dial:", err)
+	}
+	ws.Close()
+}
+
 func TestDialBadScheme(t *testing.T) {
 	s := newServer(t)
 	defer s.Close()