go-sdk/mcp/transport_test.go
Rob Findley 82be9ce756 internal/jsonrpc2_v2: refactor Reader/Writer and remove event dependency
In preparation for extracting the jsonrpc2_v2 package, update the Reader
and Writer interfaces to not report byte counts, and remove the
internal/event dependency.

Change-Id: Ie6b373736fe297c8427cd4ece7ddad33e275864b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/678956
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
2025-06-10 07:56:09 -07:00

62 lines
1.8 KiB
Go

// Copyright 2025 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 mcp
import (
"context"
"io"
"testing"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
// BatchSize causes a transport to collect n requests or notifications before
// sending a batch on the wire (responses are always sent in isolation).
//
// Exported for testing in the mcp_test package.
func BatchSize(t Transport, n int) {
if st, ok := t.(*ioStream); ok {
st.outgoingBatch = make([]jsonrpc2.Message, 0, n)
}
}
func TestBatchFraming(t *testing.T) {
// This test checks that the ndjsonFramer can read and write JSON batches.
//
// The framer is configured to write a batch size of 2, and we confirm that
// nothing is sent over the wire until the second write, at which point both
// messages become available.
ctx := context.Background()
r, w := io.Pipe()
tport := newIOStream(rwc{r, w})
tport.outgoingBatch = make([]jsonrpc2.Message, 0, 2)
// Read the two messages into a channel, for easy testing later.
read := make(chan jsonrpc2.Message)
go func() {
for range 2 {
msg, _ := tport.Read(ctx)
read <- msg
}
}()
// The first write should not yet be observed by the reader.
tport.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(1), Method: "test"})
select {
case got := <-read:
t.Fatalf("after one write, got message %v", got)
default:
}
// ...but the second write causes both messages to be observed.
tport.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(2), Method: "test"})
for _, want := range []int64{1, 2} {
got := <-read
if got := got.(*jsonrpc2.Request).ID.Raw(); got != want {
t.Errorf("got message #%d, want #%d", got, want)
}
}
}