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
1 changed files with 193 additions and 72 deletions
|
|
@ -1035,77 +1035,6 @@ func mustMarshal(v any) json.RawMessage {
|
|||
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) {
|
||||
tests := []struct {
|
||||
sid StreamID
|
||||
|
|
|
|||
Loading…
Reference in a new issue