mcp: refactor the streamable client test to be more flexible
Use a fake streamable server to facilitate testing client behavior. For this commit, just update the existing test (moved to a new file for isolation). Subsequent CLs will add more tests. Improve one client error message that occurred while debuging tests. For #393
This commit is contained in:
parent
30261728f7
commit
a4313f92bf
3 changed files with 193 additions and 72 deletions
|
|
@ -1234,7 +1234,14 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
|
||||||
|
|
||||||
default:
|
default:
|
||||||
resp.Body.Close()
|
resp.Body.Close()
|
||||||
return fmt.Errorf("unsupported content type %q", ct)
|
switch msg := msg.(type) {
|
||||||
|
case *jsonrpc.Request:
|
||||||
|
return fmt.Errorf("unsupported content type %q when sending %q (status: %d)", ct, msg.Method, resp.StatusCode)
|
||||||
|
case *jsonrpc.Response:
|
||||||
|
return fmt.Errorf("unsupported content type %q when sending jsonrpc response #%d (status: %d)", ct, msg.ID, resp.StatusCode)
|
||||||
|
default:
|
||||||
|
panic("unreachable")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
185
mcp/streamable_client_test.go
Normal file
185
mcp/streamable_client_test.go
Normal file
|
|
@ -0,0 +1,185 @@
|
||||||
|
// 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"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/google/go-cmp/cmp"
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
||||||
|
)
|
||||||
|
|
||||||
|
type streamableRequestKey struct {
|
||||||
|
httpMethod string // http method
|
||||||
|
sessionID string // session ID header
|
||||||
|
jsonrpcMethod string // jsonrpc method, or "" for non-requests
|
||||||
|
}
|
||||||
|
|
||||||
|
type header map[string]string
|
||||||
|
|
||||||
|
type streamableResponse struct {
|
||||||
|
header header
|
||||||
|
status int // or http.StatusOK
|
||||||
|
body string // or ""
|
||||||
|
optional bool // if set, request need not be sent
|
||||||
|
wantProtocolVersion string // if "", unchecked
|
||||||
|
callback func() // if set, called after the request is handled
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeResponses map[streamableRequestKey]*streamableResponse
|
||||||
|
|
||||||
|
type fakeStreamableServer struct {
|
||||||
|
t *testing.T
|
||||||
|
responses fakeResponses
|
||||||
|
|
||||||
|
callMu sync.Mutex
|
||||||
|
calls map[streamableRequestKey]int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeStreamableServer) missingRequests() []streamableRequestKey {
|
||||||
|
s.callMu.Lock()
|
||||||
|
defer s.callMu.Unlock()
|
||||||
|
|
||||||
|
var unused []streamableRequestKey
|
||||||
|
for k, resp := range s.responses {
|
||||||
|
if s.calls[k] == 0 && !resp.optional {
|
||||||
|
unused = append(unused, k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return unused
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
||||||
|
key := streamableRequestKey{
|
||||||
|
httpMethod: req.Method,
|
||||||
|
sessionID: req.Header.Get(sessionIDHeader),
|
||||||
|
}
|
||||||
|
if req.Method == http.MethodPost {
|
||||||
|
body, err := io.ReadAll(req.Body)
|
||||||
|
if err != nil {
|
||||||
|
s.t.Errorf("failed to read body: %v", err)
|
||||||
|
http.Error(w, "failed to read body", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
msg, err := jsonrpc.DecodeMessage(body)
|
||||||
|
if err != nil {
|
||||||
|
s.t.Errorf("invalid body: %v", err)
|
||||||
|
http.Error(w, "invalid body", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r, ok := msg.(*jsonrpc.Request); ok {
|
||||||
|
key.jsonrpcMethod = r.Method
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
s.callMu.Lock()
|
||||||
|
if s.calls == nil {
|
||||||
|
s.calls = make(map[streamableRequestKey]int)
|
||||||
|
}
|
||||||
|
s.calls[key]++
|
||||||
|
s.callMu.Unlock()
|
||||||
|
|
||||||
|
resp, ok := s.responses[key]
|
||||||
|
if !ok {
|
||||||
|
s.t.Errorf("missing response for %v", key)
|
||||||
|
http.Error(w, "no response", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if resp.callback != nil {
|
||||||
|
defer resp.callback()
|
||||||
|
}
|
||||||
|
for k, v := range resp.header {
|
||||||
|
w.Header().Set(k, v)
|
||||||
|
}
|
||||||
|
status := resp.status
|
||||||
|
if status == 0 {
|
||||||
|
status = http.StatusOK
|
||||||
|
}
|
||||||
|
w.WriteHeader(status)
|
||||||
|
|
||||||
|
if v := req.Header.Get(protocolVersionHeader); v != resp.wantProtocolVersion && resp.wantProtocolVersion != "" {
|
||||||
|
s.t.Errorf("%v: bad protocol version header: got %q, want %q", key, v, resp.wantProtocolVersion)
|
||||||
|
}
|
||||||
|
w.Write([]byte(resp.body))
|
||||||
|
}
|
||||||
|
|
||||||
|
var (
|
||||||
|
initResult = &InitializeResult{
|
||||||
|
Capabilities: &ServerCapabilities{
|
||||||
|
Completions: &CompletionCapabilities{},
|
||||||
|
Logging: &LoggingCapabilities{},
|
||||||
|
Tools: &ToolCapabilities{ListChanged: true},
|
||||||
|
},
|
||||||
|
ProtocolVersion: latestProtocolVersion,
|
||||||
|
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
|
||||||
|
}
|
||||||
|
initResp = resp(1, initResult, nil)
|
||||||
|
)
|
||||||
|
|
||||||
|
func jsonBody(t *testing.T, msg jsonrpc2.Message) string {
|
||||||
|
data, err := jsonrpc2.EncodeMessage(msg)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("encoding failed: %v", err)
|
||||||
|
}
|
||||||
|
return string(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStreamableClientTransportLifecycle(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// The lifecycle test verifies various behavior of the streamable client
|
||||||
|
// initialization:
|
||||||
|
// - check that it can handle application/json responses
|
||||||
|
// - check that it sends the negotiated protocol version
|
||||||
|
fake := &fakeStreamableServer{
|
||||||
|
t: t,
|
||||||
|
responses: fakeResponses{
|
||||||
|
{"POST", "", methodInitialize}: {
|
||||||
|
header: header{
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
sessionIDHeader: "123",
|
||||||
|
},
|
||||||
|
body: jsonBody(t, initResp),
|
||||||
|
},
|
||||||
|
{"POST", "123", notificationInitialized}: {
|
||||||
|
status: http.StatusAccepted,
|
||||||
|
wantProtocolVersion: latestProtocolVersion,
|
||||||
|
},
|
||||||
|
{"GET", "123", ""}: {
|
||||||
|
header: header{
|
||||||
|
"Content-Type": "text/event-stream",
|
||||||
|
},
|
||||||
|
optional: true,
|
||||||
|
wantProtocolVersion: latestProtocolVersion,
|
||||||
|
},
|
||||||
|
{"DELETE", "123", ""}: {},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
httpServer := httptest.NewServer(fake)
|
||||||
|
defer httpServer.Close()
|
||||||
|
|
||||||
|
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
||||||
|
client := NewClient(testImpl, nil)
|
||||||
|
session, err := client.Connect(ctx, transport, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("client.Connect() failed: %v", err)
|
||||||
|
}
|
||||||
|
if err := session.Close(); err != nil {
|
||||||
|
t.Errorf("closing session: %v", err)
|
||||||
|
}
|
||||||
|
if missing := fake.missingRequests(); len(missing) > 0 {
|
||||||
|
t.Errorf("did not receive expected requests: %v", missing)
|
||||||
|
}
|
||||||
|
if diff := cmp.Diff(initResult, session.state.InitializeResult); diff != "" {
|
||||||
|
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -1035,77 +1035,6 @@ func mustMarshal(v any) json.RawMessage {
|
||||||
return data
|
return data
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestStreamableClientTransport(t *testing.T) {
|
|
||||||
// This test verifies various behavior of the streamable client transport:
|
|
||||||
// - check that it can handle application/json responses
|
|
||||||
// - check that it sends the negotiated protocol version
|
|
||||||
//
|
|
||||||
// TODO(rfindley): make this test more comprehensive, similar to
|
|
||||||
// [TestStreamableServerTransport].
|
|
||||||
ctx := context.Background()
|
|
||||||
resp := func(id int64, result any, err error) *jsonrpc.Response {
|
|
||||||
return &jsonrpc.Response{
|
|
||||||
ID: jsonrpc2.Int64ID(id),
|
|
||||||
Result: mustMarshal(result),
|
|
||||||
Error: err,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
initResult := &InitializeResult{
|
|
||||||
Capabilities: &ServerCapabilities{
|
|
||||||
Completions: &CompletionCapabilities{},
|
|
||||||
Logging: &LoggingCapabilities{},
|
|
||||||
Tools: &ToolCapabilities{ListChanged: true},
|
|
||||||
},
|
|
||||||
ProtocolVersion: latestProtocolVersion,
|
|
||||||
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
|
|
||||||
}
|
|
||||||
initResp := resp(1, initResult, nil)
|
|
||||||
|
|
||||||
var reqN atomic.Int32 // request count
|
|
||||||
serverHandler := func(w http.ResponseWriter, r *http.Request) {
|
|
||||||
rN := reqN.Add(1)
|
|
||||||
|
|
||||||
// TODO(rfindley): if the status code is NoContent or Accepted, we should
|
|
||||||
// probably be tolerant of when the content type is not application/json.
|
|
||||||
w.Header().Set("Content-Type", "application/json")
|
|
||||||
if rN == 1 {
|
|
||||||
data, err := jsonrpc2.EncodeMessage(initResp)
|
|
||||||
if err != nil {
|
|
||||||
t.Errorf("encoding failed: %v", err)
|
|
||||||
}
|
|
||||||
w.Header().Set("Mcp-Session-Id", "123")
|
|
||||||
w.Write(data)
|
|
||||||
} else {
|
|
||||||
if v := r.Header.Get(protocolVersionHeader); v != latestProtocolVersion {
|
|
||||||
t.Errorf("bad protocol version header: got %q, want %q", v, latestProtocolVersion)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
httpServer := httptest.NewServer(http.HandlerFunc(serverHandler))
|
|
||||||
defer httpServer.Close()
|
|
||||||
|
|
||||||
transport := &StreamableClientTransport{Endpoint: httpServer.URL}
|
|
||||||
client := NewClient(testImpl, nil)
|
|
||||||
session, err := client.Connect(ctx, transport, nil)
|
|
||||||
if err != nil {
|
|
||||||
t.Fatalf("client.Connect() failed: %v", err)
|
|
||||||
}
|
|
||||||
if err := session.Close(); err != nil {
|
|
||||||
t.Errorf("closing session: %v", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if got, want := reqN.Load(), int32(3); got < want {
|
|
||||||
// Expect at least 3 requests: initialize, initialized, and DELETE.
|
|
||||||
// We may or may not observe the GET, depending on timing.
|
|
||||||
t.Errorf("unexpected number of requests: got %d, want at least %d", got, want)
|
|
||||||
}
|
|
||||||
|
|
||||||
if diff := cmp.Diff(initResult, session.state.InitializeResult); diff != "" {
|
|
||||||
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestEventID(t *testing.T) {
|
func TestEventID(t *testing.T) {
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
sid StreamID
|
sid StreamID
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue