From 236b9e00cd6bc438c8fa59fc9574dce02cf3f737 Mon Sep 17 00:00:00 2001 From: Sam Herrmann Date: Mon, 6 Feb 2023 17:04:43 -0500 Subject: [PATCH] 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 | 155 +++++++++++++++++++++-------------------------- 1 file changed, 69 insertions(+), 86 deletions(-) diff --git a/jsonrpc2_test.go b/jsonrpc2_test.go index a8dec6e..46b490a 100644 --- a/jsonrpc2_test.go +++ b/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 { - read, write func(p []byte) (n int, err error) +func TestConn_DisconnectNotify(t *testing.T) { + + 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) { - return x.read(p) -} - -func (x readWriteCloser) Write(p []byte) (n int, err error) { - return x.write(p) -} - -func (readWriteCloser) Close() error { return nil } - -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() { +func TestConn_Close(t *testing.T) { + t.Run("waiting for response", func(t *testing.T) { + _, connB := net.Pipe() + c := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewPlainObjectStream(connB), nil) + 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) } - close(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) - } + assertDisconnect(t, c, connB) + <-done + }) } 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) + } +}