jsonrpc: Test channel closing on when connection drops

lib/jsonrpc/rpc_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"net/http/httptest"
 	"strconv"
 	"strings"
@@ -377,6 +378,62 @@ func TestChan(t *testing.T) {
 	require.Equal(t, false, ok)
 }
 
+
+func TestChanServerClose(t *testing.T) {
+	var client struct {
+		Sub func(context.Context, int, int) (<-chan int, error)
+	}
+
+	serverHandler := &ChanHandler{
+		wait: make(chan struct{}, 5),
+	}
+
+	rpcServer := NewServer()
+	rpcServer.Register("ChanHandler", serverHandler)
+
+	tctx, tcancel := context.WithCancel(context.Background())
+
+	testServ := httptest.NewServer(rpcServer)
+	testServ.Config.ConnContext = func(ctx context.Context, c net.Conn) context.Context {
+		return tctx
+	}
+
+	closer, err := NewClient("ws://"+testServ.Listener.Addr().String(), "ChanHandler", &client, nil)
+	require.NoError(t, err)
+
+	defer closer()
+
+	serverHandler.wait <- struct{}{}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// sub
+
+	sub, err := client.Sub(ctx, 2, -1)
+	require.NoError(t, err)
+
+	// recv one
+
+	require.Equal(t, 2, <-sub)
+
+	// make sure we're blocked
+
+	select {
+	case <-time.After(200 * time.Millisecond):
+	case <-sub:
+		t.Fatal("didn't expect to get anything from sub")
+	}
+
+	// close server
+
+	tcancel()
+	testServ.Close()
+
+	_, ok := <-sub
+	require.Equal(t, false, ok)
+}
+
 func TestControlChanDeadlock(t *testing.T) {
 	for r := 0; r < 20; r++ {
 		testControlChanDeadlock(t)