internal/jsonrpc2_v2: an updated jsonrpc2 library

Change-Id: I609173baa6842d33068a7e9596d54f03d89c5401
Reviewed-on: https://go-review.googlesource.com/c/tools/+/292169
Run-TryBot: Ian Cottrell <iancottrell@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Trust: Ian Cottrell <iancottrell@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

internal/jsonrpc2/serve_test.go

@@ -0,0 +1,144 @@
+// Copyright 2020 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package jsonrpc2_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
+	"golang.org/x/tools/internal/stack/stacktest"
+)
+
+func TestIdleTimeout(t *testing.T) {
+	stacktest.NoLeak(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	listener, err := jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listener.Close()
+
+	server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{},
+		jsonrpc2.ServeOptions{
+			IdleTimeout: 100 * time.Millisecond,
+		})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	connect := func() *jsonrpc2.Connection {
+		client, err := jsonrpc2.Dial(ctx,
+			listener.Dialer(),
+			jsonrpc2.ConnectionOptions{})
+		if err != nil {
+			t.Fatal(err)
+		}
+		return client
+	}
+	// Exercise some connection/disconnection patterns, and then assert that when
+	// our timer fires, the server exits.
+	conn1 := connect()
+	conn2 := connect()
+	if err := conn1.Close(); err != nil {
+		t.Fatalf("conn1.Close failed with error: %v", err)
+	}
+	if err := conn2.Close(); err != nil {
+		t.Fatalf("conn2.Close failed with error: %v", err)
+	}
+	conn3 := connect()
+	if err := conn3.Close(); err != nil {
+		t.Fatalf("conn3.Close failed with error: %v", err)
+	}
+
+	serverError := server.Wait()
+
+	if !errors.Is(serverError, jsonrpc2.ErrIdleTimeout) {
+		t.Errorf("run() returned error %v, want %v", serverError, jsonrpc2.ErrIdleTimeout)
+	}
+}
+
+type msg struct {
+	Msg string
+}
+
+type fakeHandler struct{}
+
+func (fakeHandler) Handle(ctx context.Context, req *jsonrpc2.Request) (interface{}, error) {
+	switch req.Method {
+	case "ping":
+		return &msg{"pong"}, nil
+	default:
+		return nil, jsonrpc2.ErrNotHandled
+	}
+}
+
+func TestServe(t *testing.T) {
+	stacktest.NoLeak(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	tests := []struct {
+		name    string
+		factory func(context.Context) (jsonrpc2.Listener, error)
+	}{
+		{"tcp", func(ctx context.Context) (jsonrpc2.Listener, error) {
+			return jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
+		}},
+		{"pipe", func(ctx context.Context) (jsonrpc2.Listener, error) {
+			return jsonrpc2.NetPipe(ctx)
+		}},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			fake, err := test.factory(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			conn, shutdown, err := newFake(ctx, fake)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer shutdown(ctx)
+			var got msg
+			if err := conn.Call(ctx, "ping", &msg{"ting"}).Await(ctx, &got); err != nil {
+				t.Fatal(err)
+			}
+			if want := "pong"; got.Msg != want {
+				t.Errorf("conn.Call(...): returned %q, want %q", got, want)
+			}
+		})
+	}
+}
+
+func newFake(ctx context.Context, l jsonrpc2.Listener) (*jsonrpc2.Connection, func(context.Context), error) {
+	server, err := jsonrpc2.Serve(ctx, l, jsonrpc2.ConnectionOptions{
+		Handler: fakeHandler{},
+	}, jsonrpc2.ServeOptions{
+		IdleTimeout: 100 * time.Millisecond,
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
+	client, err := jsonrpc2.Dial(ctx,
+		l.Dialer(),
+		jsonrpc2.ConnectionOptions{
+			Handler: fakeHandler{},
+		})
+	if err != nil {
+		return nil, nil, err
+	}
+	return client, func(ctx context.Context) {
+		l.Close()
+		client.Close()
+		server.Wait()
+	}, nil
+}