jsonrpc2/conn_test.go

package jsonrpc2_test

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/sourcegraph/jsonrpc2"
)

func TestConn(t *testing.T) {

	t.Run("closes when context is done", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())

		connA, connB := Pipe(ctx, noopHandler{}, noopHandler{})
		defer connA.Close()
		defer connB.Close()

		cancel()
		<-connA.DisconnectNotify()

		got := connA.Close()
		want := jsonrpc2.ErrClosed
		if got != want {
			t.Fatalf("got %v, want %v", got, want)
		}
	})

	t.Run("cancels context when closed", func(t *testing.T) {
		ctxCanceled := make(chan struct{})

		handler := handlerFunc(func(ctx context.Context, c *jsonrpc2.Conn, r *jsonrpc2.Request) {
			// Block until the context is canceled.
			<-ctx.Done()
			close(ctxCanceled)
		})

		connA, connB := Pipe(context.Background(), noopHandler{}, jsonrpc2.AsyncHandler(handler))
		defer connA.Close()
		defer connB.Close()

		// Send a notification from connA to connB to trigger connB's handler
		// function.
		if err := connA.Notify(context.Background(), "foo", nil, nil); err != nil {
			t.Fatal(err)
		}

		// Disconnect connA from connB.
		if err := connA.Close(); err != nil {
			t.Fatal(err)
		}

		select {
		case <-ctxCanceled:
			// Test passed, the handler's context was canceled.
		case <-time.After(time.Second):
			t.Fatal("context not canceled")
		}
	})
}

var paramsTests = []struct {
	sendParams interface{}
	wantParams *json.RawMessage
}{
	{
		sendParams: nil,
		wantParams: nil,
	},
	{
		sendParams: jsonNull,
		wantParams: &jsonNull,
	},
	{
		sendParams: false,
		wantParams: rawJSONMessage("false"),
	},
	{
		sendParams: 0,
		wantParams: rawJSONMessage("0"),
	},
	{
		sendParams: "",
		wantParams: rawJSONMessage(`""`),
	},
	{
		sendParams: rawJSONMessage(`{"foo":"bar"}`),
		wantParams: rawJSONMessage(`{"foo":"bar"}`),
	},
}

func TestConn_DispatchCall(t *testing.T) {
	for _, test := range paramsTests {
		t.Run(fmt.Sprintf("%s", test.sendParams), func(t *testing.T) {
			testParams(t, test.wantParams, func(c *jsonrpc2.Conn) error {
				_, err := c.DispatchCall(context.Background(), "f", test.sendParams)
				return err
			})
		})
	}
}

func TestConn_Notify(t *testing.T) {
	for _, test := range paramsTests {
		t.Run(fmt.Sprintf("%s", test.sendParams), func(t *testing.T) {
			testParams(t, test.wantParams, func(c *jsonrpc2.Conn) error {
				return c.Notify(context.Background(), "f", test.sendParams)
			})
		})
	}
}

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),
			noopHandler{},
			// Suppress log message. This connection receives an invalid JSON
			// message that causes an error to be written to the logger. We
			// don't want this expected error to appear in os.Stderr though when
			// running tests in verbose mode or when other tests fail.
			jsonrpc2.SetLogger(log.New(io.Discard, "", 0)),
		)
		connA.Write([]byte("invalid json"))
		assertDisconnect(t, c, connB)
	})
}

func TestConn_Close(t *testing.T) {
	cases := []struct {
		name string
		run  func(*testing.T, context.Context, *jsonrpc2.Conn)
	}{{
		name: "during Call",
		run: func(t *testing.T, ctx context.Context, conn *jsonrpc2.Conn) {
			ready := make(chan struct{})
			done := make(chan struct{})
			go func() {
				close(ready)
				err := conn.Call(ctx, "m", nil, nil)
				if err != jsonrpc2.ErrClosed {
					t.Errorf("got error %v, want %v", err, jsonrpc2.ErrClosed)
				}
				close(done)
			}()
			// Wait for the request to be sent before we close the connection.
			<-ready
			if err := conn.Close(); err != nil && err != jsonrpc2.ErrClosed {
				t.Error(err)
			}
			<-done
		},
	}, {
		name: "during Wait",
		run: func(t *testing.T, ctx context.Context, conn *jsonrpc2.Conn) {
			call, err := conn.DispatchCall(ctx, "m", nil, nil)
			if err != nil {
				t.Fatal(err)
			}
			if err := conn.Close(); err != nil {
				t.Fatal(err)
			}
			if err := call.Wait(ctx, nil); err != jsonrpc2.ErrClosed {
				t.Fatal(err)
			}
		},
	}, {
		name: "during Dispatch",
		run: func(t *testing.T, ctx context.Context, conn *jsonrpc2.Conn) {
			if err := conn.Close(); err != nil {
				t.Fatal(err)
			}
			if _, err := conn.DispatchCall(ctx, "m", nil, nil); err != jsonrpc2.ErrClosed {
				t.Fatal(err)
			}
		},
	}}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			connA, connB := net.Pipe()
			nodeA := jsonrpc2.NewConn(
				ctx,
				jsonrpc2.NewPlainObjectStream(connA), noopHandler{},
			)
			defer nodeA.Close()
			nodeB := jsonrpc2.NewConn(
				ctx,
				jsonrpc2.NewPlainObjectStream(connB),
				noopHandler{},
			)
			defer nodeB.Close()

			tc.run(t, ctx, nodeB)

			assertDisconnect(t, nodeB, connB)
		})
	}
}

func testParams(t *testing.T, want *json.RawMessage, fn func(c *jsonrpc2.Conn) error) {
	wg := &sync.WaitGroup{}
	handler := handlerFunc(func(ctx context.Context, conn *jsonrpc2.Conn, r *jsonrpc2.Request) {
		assertRawJSONMessage(t, r.Params, want)
		wg.Done()
	})

	connA, connB := Pipe(context.Background(), noopHandler{}, handler)
	defer connA.Close()
	defer connB.Close()

	wg.Add(1)
	if err := fn(connA); err != nil {
		t.Error(err)
	}
	wg.Wait()
}

func assertDisconnect(t *testing.T, c *jsonrpc2.Conn, conn io.Writer) {
	select {
	case <-c.DisconnectNotify():
	case <-time.After(200 * time.Millisecond):
		t.Error("no disconnect notification")
		return
	}
	// Assert that conn is closed by trying to write to it.
	_, got := conn.Write(nil)
	want := io.ErrClosedPipe
	if got != want {
		t.Errorf("got %s, want %s", got, want)
	}
}

func assertRawJSONMessage(t *testing.T, got *json.RawMessage, want *json.RawMessage) {
	// Assert pointers.
	if got == nil || want == nil {
		if got != want {
			t.Errorf("pointer: got %s, want %s", got, want)
		}
		return
	}
	{
		// If pointers are not nil, then assert values.
		got := string(*got)
		want := string(*want)
		if got != want {
			t.Errorf("value: got %q, want %q", got, want)
		}
	}
}

// Pipe returns two jsonrpc2.Conn, connected via a synchronous, in-memory, full
// duplex network connection.
func Pipe(ctx context.Context, handlerA, handlerB jsonrpc2.Handler) (connA *jsonrpc2.Conn, connB *jsonrpc2.Conn) {
	a, b := net.Pipe()
	connA = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(a), handlerA)
	connB = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(b), handlerB)
	return connA, connB
}