1
0
Fork 0
mirror of https://github.com/sourcegraph/jsonrpc2.git synced 2026-06-16 04:04:56 +02:00

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.
This commit is contained in:
Sam Herrmann 2023-02-06 17:04:43 -05:00
parent 21db451b57
commit 236b9e00cd

View file

@ -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)
}
}