Cleanup TestConn_DisconnectNotify

The purpose of this commit is to consolidate assertions that are
required for multiple tests related to `Conn.DisconnectNotify` into a
single function.

jsonrpc2_test.go

@@ -314,97 +314,66 @@ type noopHandler struct{}
 
 func (noopHandler) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {}
 
-type readWriteCloser struct {
+func TestConn_DisconnectNotify(t *testing.T) {
-	read, write func(p []byte) (n int, err error)
+
+	t.Run("EOF", func(t *testing.T) {
+		connA, connB := net.Pipe()
+		c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil)
+		// By closing connA, connB receives io.EOF
+		if err := connA.Close(); err != nil {
+			t.Error(err)
+		}
+		assertDisconnect(t, c, connB)
+	})
+
+	t.Run("Close", func(t *testing.T) {
+		_, connB := net.Pipe()
+		c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil)
+		if err := c.Close(); err != nil {
+			t.Error(err)
+		}
+		assertDisconnect(t, c, connB)
+	})
+
+	t.Run("Close async", func(t *testing.T) {
+		done := make(chan struct{})
+		_, connB := net.Pipe()
+		c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil)
+		go func() {
+			if err := c.Close(); err != nil && err != jsonrpc2.ErrClosed {
+				t.Error(err)
+			}
+			close(done)
+		}()
+		assertDisconnect(t, c, connB)
+		<-done
+	})
+
+	t.Run("protocol error", func(t *testing.T) {
+		connA, connB := net.Pipe()
+		c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil)
+		connA.Write([]byte("invalid json"))
+		assertDisconnect(t, c, connB)
+	})
 }
 
-func (x readWriteCloser) Read(p []byte) (n int, err error) {
+func TestConn_Close(t *testing.T) {
-	return x.read(p)
+	t.Run("waiting for response", func(t *testing.T) {
-}
+		_, connB := net.Pipe()
-
+		c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil)
-func (x readWriteCloser) Write(p []byte) (n int, err error) {
+		done := make(chan struct{})
-	return x.write(p)
+		go func() {
-}
+			if err := c.Call(context.Background(), "m", nil, nil); err != jsonrpc2.ErrClosed {
-
+				t.Errorf("got error %v, want %v", err, jsonrpc2.ErrClosed)
-func (readWriteCloser) Close() error { return nil }
+			}
-
+			close(done)
-func eof(p []byte) (n int, err error) {
+		}()
-	return 0, io.EOF
-}
-
-func TestConn_DisconnectNotify_EOF(t *testing.T) {
-	c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(&readWriteCloser{eof, eof}, jsonrpc2.VarintObjectCodec{}), nil)
-	select {
-	case <-c.DisconnectNotify():
-	case <-time.After(200 * time.Millisecond):
-		t.Fatal("no disconnect notification")
-	}
-}
-
-func TestConn_DisconnectNotify_Close(t *testing.T) {
-	c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(&readWriteCloser{eof, eof}, jsonrpc2.VarintObjectCodec{}), nil)
-	if err := c.Close(); err != nil {
-		t.Error(err)
-	}
-	select {
-	case <-c.DisconnectNotify():
-	case <-time.After(200 * time.Millisecond):
-		t.Fatal("no disconnect notification")
-	}
-}
-
-func TestConn_DisconnectNotify_Close_async(t *testing.T) {
-	done := make(chan struct{})
-	c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(&readWriteCloser{eof, eof}, jsonrpc2.VarintObjectCodec{}), nil)
-	go func() {
 		if err := c.Close(); err != nil && err != jsonrpc2.ErrClosed {
 			t.Error(err)
 		}
-		close(done)
+		assertDisconnect(t, c, connB)
-	}()
+		<-done
-	select {
+	})
-	case <-c.DisconnectNotify():
-	case <-time.After(200 * time.Millisecond):
-		t.Fatal("no disconnect notification")
-	}
-	<-done
-}
-
-func TestConn_Close_waitingForResponse(t *testing.T) {
-	c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(&readWriteCloser{eof, eof}, jsonrpc2.VarintObjectCodec{}), noopHandler{})
-	done := make(chan struct{})
-	go func() {
-		if err := c.Call(context.Background(), "m", nil, nil); err != jsonrpc2.ErrClosed {
-			t.Errorf("got error %v, want %v", err, jsonrpc2.ErrClosed)
-		}
-		close(done)
-	}()
-	if err := c.Close(); err != nil && err != jsonrpc2.ErrClosed {
-		t.Error(err)
-	}
-	select {
-	case <-c.DisconnectNotify():
-	case <-time.After(200 * time.Millisecond):
-		t.Fatal("no disconnect notification")
-	}
-	<-done
-}
-
-func TestConn_DisconnectNotify_protocol_error(t *testing.T) {
-	connA, connB := net.Pipe()
-	c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(connB, jsonrpc2.VarintObjectCodec{}), nil)
-	connA.Write([]byte("invalid json"))
-	select {
-	case <-c.DisconnectNotify():
-	case <-time.After(200 * time.Millisecond):
-		t.Fatal("no disconnect notification")
-	}
-	// Assert that the underlying connection is closed by trying to write to it.
-	_, got := connB.Write(nil)
-	want := io.ErrClosedPipe
-	if got != want {
-		t.Fatalf("got %q, want %q", got, want)
-	}
 }
 
 func serve(ctx context.Context, lis net.Listener, h jsonrpc2.Handler, streamMaker streamMaker, opts ...jsonrpc2.ConnOpt) error {
@@ -416,3 +385,17 @@ func serve(ctx context.Context, lis net.Listener, h jsonrpc2.Handler, streamMake
 		jsonrpc2.NewConn(ctx, streamMaker(conn), h, opts...)
 	}
 }
+
+func assertDisconnect(t *testing.T, c *jsonrpc2.Conn, conn net.Conn) {
+	select {
+	case <-c.DisconnectNotify():
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("no disconnect notification")
+	}
+	// Assert that conn is closed by trying to write to it.
+	_, got := conn.Write(nil)
+	want := io.ErrClosedPipe
+	if got != want {
+		t.Fatalf("got %q, want %q", got, want)
+	}
+}