internal/mcp/transport_test.go
Rama main · réplica de lectura del repositorio original.
package mcp
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"testing"
"time"
"github.com/sophotechlabs/spinoza/internal/api"
)
type interruptedInput struct {
delivered bool
}
func (i *interruptedInput) Read(into []byte) (int, error) {
if i.delivered {
return 0, errRefused
}
i.delivered = true
return copy(into, `{"jsonrpc":"2.0","id":1,"method":"ping"}`), nil
}
func TestStdioAnswersOneMessagePerLine(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
in := strings.NewReader(
`{"jsonrpc":"2.0","id":1,"method":"ping"}` + "\n" +
`{"jsonrpc":"2.0","id":2,"method":"ping"}` + "\n",
)
var out bytes.Buffer
if err := server.Serve(context.Background(), in, &out); err != nil {
t.Fatalf("serve: %v", err)
}
lines := strings.Split(strings.TrimSpace(out.String()), "\n")
if len(lines) != 2 {
t.Fatalf("wrote %d lines, want one per request:\n%s", len(lines), out.String())
}
for index, line := range lines {
var reply map[string]any
if err := json.Unmarshal([]byte(line), &reply); err != nil {
t.Fatalf("line %d is not JSON: %v", index, err)
}
}
}
func TestStdioSkipsBlankLinesAndSaysNothingToNotifications(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
in := strings.NewReader("\n" + `{"jsonrpc":"2.0","method":"notifications/initialized"}` + "\n\n")
var out bytes.Buffer
if err := server.Serve(context.Background(), in, &out); err != nil {
t.Fatalf("serve: %v", err)
}
if out.Len() != 0 {
t.Fatalf("wrote %q, want nothing", out.String())
}
}
func TestStdioStopsWhenTheContextIsDone(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
ctx, cancel := context.WithCancel(context.Background())
cancel()
var out bytes.Buffer
err := server.Serve(ctx, strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`+"\n"), &out)
if err == nil {
t.Fatal("serve carried on after the context was canceled")
}
}
func TestTheToolListingMarksWhatWrites(t *testing.T) {
server := serverFor(writableCluster(), Options{AllowWrite: true})
var out bytes.Buffer
if err := server.List(&out); err != nil {
t.Fatalf("list: %v", err)
}
body := out.String()
if !strings.Contains(body, "get_dashboard read") {
t.Fatalf("listing does not mark a read tool:\n%s", body)
}
if !strings.Contains(body, "manage_node write") {
t.Fatalf("listing does not mark a write tool:\n%s", body)
}
}
func TestCallingOneToolFromTheCommandLine(t *testing.T) {
cluster := &fakeCluster{spaces: api.Namespaces{Names: []string{"prod"}}}
server := serverFor(cluster, Options{})
var out bytes.Buffer
if err := server.Call(context.Background(), &out, "list_namespaces", nil); err != nil {
t.Fatalf("call: %v", err)
}
if !strings.Contains(out.String(), `"name": "prod"`) {
t.Fatalf("output = %s", out.String())
}
}
func TestCallingWithArgumentsFromTheCommandLine(t *testing.T) {
cluster := &fakeCluster{catalog: catalogOf(descriptor("", "v1", "pods", "Pod"))}
server := serverFor(cluster, Options{})
var out bytes.Buffer
err := server.Call(context.Background(), &out, "list_resources", []string{"resource=pods", "namespace=prod"})
if err != nil {
t.Fatalf("call: %v", err)
}
if cluster.lastKind.Resource != "pods" {
t.Fatalf("resource = %q", cluster.lastKind.Resource)
}
}
func TestCommandLineArgumentsCarryTheirType(t *testing.T) {
cases := []struct {
name string
pair string
key string
value any
}{
{name: "a word", pair: "resource=pods", key: "resource", value: "pods"},
{name: "a number", pair: "limit=5", key: "limit", value: float64(5)},
{name: "true", pair: "events=true", key: "events", value: true},
{name: "false", pair: "events=false", key: "events", value: false},
{name: "a value with an equals in it", pair: "query=up==1", key: "query", value: "up==1"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
args, err := parsePairs([]string{tc.pair})
if err != nil {
t.Fatalf("parse: %v", err)
}
if args[tc.key] != tc.value {
t.Fatalf("%s = %#v, want %#v", tc.key, args[tc.key], tc.value)
}
})
}
}
func TestAnArgumentWithNoEqualsIsRefused(t *testing.T) {
if _, err := parsePairs([]string{"resource"}); err == nil {
t.Fatal("a bare word was accepted as an argument")
}
}
func TestACommandLineCallRejectsAnArgumentWithNoEquals(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
err := server.Call(t.Context(), &out, "list_namespaces", []string{"namespace"})
if err == nil || !strings.Contains(err.Error(), "key=value") {
t.Fatalf("error = %v, want the malformed pair named", err)
}
if out.Len() != 0 {
t.Fatalf("output = %q, want no partial answer", out.String())
}
}
func TestCallingSomethingTheServerDoesNotServe(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
err := server.Call(context.Background(), &out, "manage_node", nil)
if err == nil || !strings.Contains(err.Error(), "read-only") {
t.Fatalf("error = %v, want the read-only reason", err)
}
}
func TestCallingAToolThatFailsReportsIt(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
if err := server.Call(context.Background(), &out, "search", nil); err == nil {
t.Fatal("a failing tool reported success")
}
if _, err := parsePairs([]string{"a=b"}); err != nil {
t.Fatalf("parse: %v", err)
}
}
func TestEveryResourceIsListedWithAURIAndAType(t *testing.T) {
for _, card := range resourceCards() {
if !strings.HasPrefix(card.URI, "cluster://") {
t.Fatalf("%s does not look like a cluster resource", card.URI)
}
if card.MimeType != jsonMime {
t.Fatalf("%s is %q, want %s", card.URI, card.MimeType, jsonMime)
}
if card.Name == "" || card.Description == "" {
t.Fatalf("%s is unlabelled", card.URI)
}
}
}
func TestReadingEachResource(t *testing.T) {
cluster := &fakeCluster{
overview: api.ClusterOverview{Version: "v1.30.0"},
graph: api.Graph{Nodes: []api.GraphNode{{ID: "a", Name: "web"}}},
events: []api.Event{{Reason: "Pulled", Message: "ok"}},
}
server := serverFor(cluster, Options{})
for _, card := range resourceCards() {
t.Run(card.URI, func(t *testing.T) {
body, err := server.resourceBody(context.Background(), card.URI)
if err != nil {
t.Fatalf("read %s: %v", card.URI, err)
}
if body == nil {
t.Fatalf("%s came back empty", card.URI)
}
})
}
}
func TestReadingAResourceNobodyServes(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
_, err := server.resourceBody(context.Background(), "cluster://secrets")
if err == nil || !strings.Contains(err.Error(), "cluster://secrets") {
t.Fatalf("error = %v, want the uri named", err)
}
}
func TestReadingAResourceOverTheProtocol(t *testing.T) {
server := serverFor(&fakeCluster{overview: api.ClusterOverview{Version: "v1.30.0"}}, Options{})
reply := ask(t, server, `{"jsonrpc":"2.0","id":1,"method":"resources/read","params":{"uri":"cluster://health"}}`)
contents := as[[]any](t, resultOf(t, reply)["contents"])
first := as[map[string]any](t, contents[0])
if first["mimeType"] != jsonMime {
t.Fatalf("mimeType = %v", first["mimeType"])
}
if !strings.Contains(as[string](t, first["text"]), "v1.30.0") {
t.Fatalf("text = %v", first["text"])
}
}
func TestReadParametersThatAreNotAnObjectAreRefused(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
reply := ask(t, server, `{"jsonrpc":"2.0","id":1,"method":"resources/read","params":[]}`)
if errorOf(t, reply)["code"] != float64(codeInvalidParams) {
t.Fatalf("code = %v", errorOf(t, reply)["code"])
}
}
func TestAskingForAResourceThatIsNotThereOverTheProtocol(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
reply := ask(t, server, `{"jsonrpc":"2.0","id":1,"method":"resources/read","params":{"uri":"cluster://nope"}}`)
if errorOf(t, reply)["code"] != float64(codeInvalidParams) {
t.Fatalf("code = %v", errorOf(t, reply)["code"])
}
}
type brokenWriter struct{}
func (brokenWriter) Write([]byte) (int, error) { return 0, errRefused }
func TestAWriteFailureIsReportedRatherThanSwallowed(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
if err := server.List(brokenWriter{}); !errors.Is(err, errRefused) {
t.Fatalf("list error = %v", err)
}
if err := server.Call(context.Background(), brokenWriter{}, "list_namespaces", nil); !errors.Is(err, errRefused) {
t.Fatalf("call error = %v", err)
}
err := server.Serve(context.Background(), strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`+"\n"), brokenWriter{})
if !errors.Is(err, errRefused) {
t.Fatalf("serve error = %v", err)
}
}
func TestAWriteFailureWhileFillingTheBufferIsReported(t *testing.T) {
sink := &writer{out: bufio.NewWriter(brokenWriter{})}
sink.send(bytes.Repeat([]byte("x"), 8192))
if !errors.Is(sink.err(), errRefused) {
t.Fatalf("write error = %v, want the underlying failure", sink.err())
}
}
func TestAStdioInputFailureIsReturnedAfterItsLastMessage(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
err := server.Serve(t.Context(), &interruptedInput{}, &out)
if !errors.Is(err, errRefused) {
t.Fatalf("serve error = %v, want the input failure", err)
}
if !strings.Contains(out.String(), `"jsonrpc":"2.0"`) {
t.Fatalf("output = %q, want the complete final request answered", out.String())
}
}
func TestListingToolsOverTheProtocol(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
result := resultOf(t, ask(t, server, `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))
tools := as[[]any](t, result["tools"])
if len(tools) == 0 {
t.Fatal("tools/list came back empty")
}
first := as[map[string]any](t, tools[0])
if first["name"] == nil || first["inputSchema"] == nil {
t.Fatalf("card = %v", first)
}
}
func TestListingResourcesOverTheProtocol(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
result := resultOf(t, ask(t, server, `{"jsonrpc":"2.0","id":1,"method":"resources/list"}`))
if len(as[[]any](t, result["resources"])) != len(resourceCards()) {
t.Fatalf("resources = %v", result["resources"])
}
}
func TestASuccessfulCallComesBackAsTextContent(t *testing.T) {
cluster := &fakeCluster{spaces: api.Namespaces{Names: []string{"prod"}}}
server := serverFor(cluster, Options{})
result := resultOf(t, ask(t, server,
`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"list_namespaces","arguments":{}}}`))
if result["isError"] != nil {
t.Fatalf("result = %v, want no error flag", result)
}
contents := as[[]any](t, result["content"])
body := as[map[string]any](t, contents[0])
if body["type"] != "text" {
t.Fatalf("content type = %v", body["type"])
}
if !strings.Contains(as[string](t, body["text"]), "prod") {
t.Fatalf("text = %v", body["text"])
}
}
func TestParsingTheCommandLine(t *testing.T) {
var out bytes.Buffer
opts, err := Parse([]string{
"-context", "p-mk1", "-allow-write", "-unsafe-raw-output", "-log-lines", "50",
"call", "get_dashboard",
}, &out)
if err != nil {
t.Fatalf("parse: %v", err)
}
if opts.Context != "p-mk1" {
t.Fatalf("context = %q", opts.Context)
}
if !opts.AllowWrite {
t.Fatal("allow-write was not read")
}
if !opts.UnsafeRawOutput {
t.Fatal("unsafe-raw-output was not read")
}
if opts.LogLines != 50 {
t.Fatalf("logLines = %d", opts.LogLines)
}
if len(opts.Args) != 2 || opts.Args[0] != "call" {
t.Fatalf("args = %v", opts.Args)
}
}
func TestTheDefaultsWhenNothingIsAsked(t *testing.T) {
var out bytes.Buffer
opts, err := Parse(nil, &out)
if err != nil {
t.Fatalf("parse: %v", err)
}
if opts.LogLines != defaultLogLines {
t.Fatalf("logLines = %d, want %d", opts.LogLines, defaultLogLines)
}
if opts.AllowWrite {
t.Fatal("writes are on by default; they must be asked for")
}
if opts.UnsafeRawOutput {
t.Fatal("raw logs and values are on by default")
}
if opts.SyncWait == 0 {
t.Fatal("no sync timeout was set, so a cold cache would wait forever")
}
if len(opts.Args) != 0 {
t.Fatalf("args = %v, want none", opts.Args)
}
}
func TestAskingForHelpIsNotAFailure(t *testing.T) {
var out bytes.Buffer
_, err := Parse([]string{"-h"}, &out)
if !errors.Is(err, ErrHelp) {
t.Fatalf("error = %v, want ErrHelp", err)
}
if !strings.Contains(out.String(), "spinoza-mcp") {
t.Fatalf("usage = %q, want it to name the command", out.String())
}
}
func TestAFlagNobodyKnowsIsRefused(t *testing.T) {
var out bytes.Buffer
if _, err := Parse([]string{"-nonsense"}, &out); err == nil || errors.Is(err, ErrHelp) {
t.Fatalf("error = %v, want a real failure", err)
}
}
func TestDispatchPicksTheModeFromTheArguments(t *testing.T) {
cluster := &fakeCluster{spaces: api.Namespaces{Names: []string{"prod"}}}
server := serverFor(cluster, Options{})
cases := []struct {
name string
args []string
in string
wants string
fails string
}{
{name: "no arguments serves the protocol", args: nil, in: `{"jsonrpc":"2.0","id":1,"method":"ping"}` + "\n", wants: `"id":1`},
{name: "tools lists them", args: []string{"tools"}, wants: "list_namespaces"},
{name: "call runs one", args: []string{"call", "list_namespaces"}, wants: "prod"},
{name: "call with no name", args: []string{"call"}, fails: "call needs a tool name"},
{name: "a command nobody knows", args: []string{"drain"}, fails: `unknown command "drain"`},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
var out bytes.Buffer
err := server.Dispatch(context.Background(), Settings{Args: tc.args}, strings.NewReader(tc.in), &out)
if tc.fails != "" {
if err == nil || !strings.Contains(err.Error(), tc.fails) {
t.Fatalf("error = %v, want %q", err, tc.fails)
}
return
}
if err != nil {
t.Fatalf("dispatch: %v", err)
}
if !strings.Contains(out.String(), tc.wants) {
t.Fatalf("output = %q, want it to contain %q", out.String(), tc.wants)
}
})
}
}
func TestPrometheusIsSkippedWhenTheSpecIsNonsense(t *testing.T) {
if PromFor(api.ContextRef{Name: "nowhere"}, Settings{PromSpec: "not a target"}) != nil {
t.Fatal("a Prometheus client was built from a spec that does not parse")
}
}
func TestPrometheusIsSkippedWhenTheClusterCannotBeLoaded(t *testing.T) {
opts := Settings{
Kubeconfig: filepath.Join(t.TempDir(), "missing"),
PromSpec: "monitoring/prometheus:9090",
}
if PromFor(api.ContextRef{Name: "test"}, opts) != nil {
t.Fatal("a Prometheus client was built without a readable cluster")
}
}
func TestPrometheusUsesTheSelectedCluster(t *testing.T) {
path := filepath.Join(t.TempDir(), "config")
body := `apiVersion: v1
kind: Config
current-context: test
clusters:
- name: local
cluster:
server: https://127.0.0.1:6443
insecure-skip-tls-verify: true
contexts:
- name: test
context:
cluster: local
namespace: default
user: test
users:
- name: test
user:
token: token
`
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
t.Fatalf("write kubeconfig: %v", err)
}
opts := Settings{Kubeconfig: path, PromSpec: "monitoring/prometheus:9090"}
if PromFor(api.ContextRef{Name: "test"}, opts) == nil {
t.Fatal("a valid cluster and Prometheus target produced no client")
}
}
func TestAMessageTooLargeIsAnsweredRatherThanEndingTheSession(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
huge := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"search","arguments":{"query":"` +
strings.Repeat("x", maxMessage) + `"}}}`
in := strings.NewReader(huge + "\n" + `{"jsonrpc":"2.0","id":2,"method":"ping"}` + "\n")
var out bytes.Buffer
if err := server.Serve(context.Background(), in, &out); err != nil {
t.Fatalf("an oversized message ended the session: %v", err)
}
body := out.String()
if !strings.Contains(body, "larger than this server accepts") {
t.Fatalf("the oversized message drew no answer:\n%s", body)
}
if !strings.Contains(body, `"id":2`) {
t.Fatalf("the message after the oversized one was never served:\n%s", body)
}
}
func TestReadingAnOversizedMessageRetainsOnlyTheLimit(t *testing.T) {
oversized := strings.Repeat("x", maxMessage+64*1024)
reader := bufio.NewReaderSize(strings.NewReader(oversized+"\nnext\n"), 64*1024)
first, err := readMessage(reader)
if err != nil {
t.Fatalf("first message: %v", err)
}
if len(first) != maxMessage+2 {
t.Fatalf("retained = %d bytes, want %d", len(first), maxMessage+2)
}
second, secondErr := readMessage(reader)
if secondErr != nil {
t.Fatalf("second message: %v", secondErr)
}
if string(second) != "next" {
t.Fatalf("second message = %q, want the request after the oversized line", second)
}
}
func TestALineLongerThanTheBufferIsStillOneMessage(t *testing.T) {
padding := strings.Repeat("y", 200*1024)
in := strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping","params":{"pad":"` + padding + `"}}` + "\n")
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
if err := server.Serve(context.Background(), in, &out); err != nil {
t.Fatalf("serve: %v", err)
}
if !strings.Contains(out.String(), `"id":1`) {
t.Fatalf("a message larger than the read buffer was not reassembled:\n%s", out.String())
}
if strings.Count(strings.TrimSpace(out.String()), "\n") != 0 {
t.Fatalf("one message drew more than one reply:\n%s", out.String())
}
}
func TestALineWithNoNewlineAtTheEndIsStillServed(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
var out bytes.Buffer
if err := server.Serve(context.Background(), strings.NewReader(`{"jsonrpc":"2.0","id":9,"method":"ping"}`), &out); err != nil {
t.Fatalf("serve: %v", err)
}
if !strings.Contains(out.String(), `"id":9`) {
t.Fatalf("the last message was dropped for want of a newline:\n%s", out.String())
}
}
func TestASlowToolDoesNotHoldUpTheOnesBehindIt(t *testing.T) {
release := make(chan struct{})
releaseAll := closeOnce(release)
defer releaseAll()
cluster := &blockingCluster{fakeCluster: &fakeCluster{}, release: release}
server := serverFor(cluster, Options{})
in := strings.NewReader(
`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"list_namespaces","arguments":{}}}` + "\n" +
`{"jsonrpc":"2.0","id":2,"method":"ping"}` + "\n",
)
out := lockedBuffer{written: make(chan struct{}, 4)}
done := make(chan error, 1)
go func() { done <- server.Serve(context.Background(), in, &out) }()
for !strings.Contains(out.read(), `"id":2`) {
select {
case <-out.written:
case <-time.After(5 * time.Second):
t.Fatal("the ping never came back while a slow tool was running")
}
}
releaseAll()
if err := <-done; err != nil {
t.Fatalf("serve: %v", err)
}
if !strings.Contains(out.read(), `"id":1`) {
t.Fatalf("the slow tool never answered:\n%s", out.read())
}
}
type blockingCluster struct {
*fakeCluster
release chan struct{}
entered chan struct{}
}
type panickingCluster struct {
*fakeCluster
}
func (p *panickingCluster) Namespaces(context.Context) api.Namespaces {
panic("broken namespace reader")
}
func TestStdioTurnsAHandlerPanicIntoAnInternalError(t *testing.T) {
server := serverFor(&panickingCluster{fakeCluster: &fakeCluster{}}, Options{})
var out bytes.Buffer
err := server.Serve(context.Background(), strings.NewReader(namespaceRequests(1)), &out)
if err != nil {
t.Fatalf("serve: %v", err)
}
if !strings.Contains(out.String(), `"id":1`) {
t.Fatalf("reply = %s, want the request id", out.String())
}
if !strings.Contains(out.String(), `"code":-32603`) {
t.Fatalf("reply = %s, want an internal error", out.String())
}
}
func (b *blockingCluster) Namespaces(context.Context) api.Namespaces {
if b.entered != nil {
b.entered <- struct{}{}
}
<-b.release
return api.Namespaces{Names: []string{"prod"}}
}
func namespaceRequests(count int) string {
var in strings.Builder
for id := 1; id <= count; id++ {
_, _ = fmt.Fprintf(
&in,
`{"jsonrpc":"2.0","id":%d,"method":"tools/call","params":{"name":"list_namespaces","arguments":{}}}`+"\n",
id,
)
}
return in.String()
}
func waitForEntries(t *testing.T, entered <-chan struct{}, count int) {
t.Helper()
timer := time.NewTimer(5 * time.Second)
defer timer.Stop()
for received := range count {
select {
case <-entered:
case <-timer.C:
t.Fatalf("calls entered = %d, want %d", received, count)
}
}
}
func closeOnce(ch chan struct{}) func() {
var once sync.Once
return func() {
once.Do(func() {
close(ch)
})
}
}
func TestStdioCapsConcurrentCalls(t *testing.T) {
release := make(chan struct{})
releaseAll := closeOnce(release)
entered := make(chan struct{}, atOnce+1)
cluster := &blockingCluster{fakeCluster: &fakeCluster{}, release: release, entered: entered}
server := serverFor(cluster, Options{})
var out lockedBuffer
done := make(chan error, 1)
finished := make(chan struct{})
t.Cleanup(func() {
releaseAll()
select {
case <-finished:
case <-time.After(5 * time.Second):
t.Error("serve did not stop during cleanup")
}
})
go func() {
defer close(finished)
done <- server.Serve(context.Background(), strings.NewReader(namespaceRequests(atOnce+1)), &out)
}()
waitForEntries(t, entered, atOnce)
select {
case <-entered:
t.Fatalf("more than %d calls ran at once", atOnce)
default:
}
releaseAll()
if err := <-done; err != nil {
t.Fatalf("serve: %v", err)
}
if replies := strings.Count(strings.TrimSpace(out.read()), "\n") + 1; replies != atOnce+1 {
t.Fatalf("replies = %d, want %d", replies, atOnce+1)
}
}
func TestCancellationWaitsForRunningCalls(t *testing.T) {
release := make(chan struct{})
releaseAll := closeOnce(release)
entered := make(chan struct{}, atOnce+1)
cluster := &blockingCluster{fakeCluster: &fakeCluster{}, release: release, entered: entered}
server := serverFor(cluster, Options{})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
done := make(chan error, 1)
finished := make(chan struct{})
t.Cleanup(func() {
releaseAll()
select {
case <-finished:
case <-time.After(5 * time.Second):
t.Error("serve did not stop during cleanup")
}
})
go func() {
defer close(finished)
done <- server.Serve(ctx, strings.NewReader(namespaceRequests(atOnce+1)), &lockedBuffer{})
}()
waitForEntries(t, entered, atOnce)
cancel()
select {
case err := <-done:
t.Fatalf("serve returned before its workers: %v", err)
default:
}
releaseAll()
select {
case err := <-done:
if !errors.Is(err, context.Canceled) {
t.Fatalf("serve error = %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("serve did not finish after its workers")
}
}
type yieldingWriter struct {
body bytes.Buffer
}
func (w *yieldingWriter) Write(p []byte) (int, error) {
middle := len(p) / 2
first, _ := w.body.Write(p[:middle])
runtime.Gosched()
second, _ := w.body.Write(p[middle:])
return first + second, nil
}
func TestConcurrentRepliesStayIndivisible(t *testing.T) {
server := serverFor(&fakeCluster{}, Options{})
const calls = 64
var in strings.Builder
for id := 1; id <= calls; id++ {
_, _ = fmt.Fprintf(&in, `{"jsonrpc":"2.0","id":%d,"method":"ping"}`+"\n", id)
}
var out yieldingWriter
if err := server.Serve(context.Background(), strings.NewReader(in.String()), &out); err != nil {
t.Fatalf("serve: %v", err)
}
lines := strings.Split(strings.TrimSpace(out.body.String()), "\n")
if len(lines) != calls {
t.Fatalf("replies = %d, want %d", len(lines), calls)
}
ids := map[float64]bool{}
for _, line := range lines {
var reply map[string]any
if err := json.Unmarshal([]byte(line), &reply); err != nil {
t.Fatalf("interleaved reply %q: %v", line, err)
}
id, ok := reply["id"].(float64)
if !ok {
t.Fatalf("reply id = %#v, want a number", reply["id"])
}
ids[id] = true
}
if len(ids) != calls {
t.Fatalf("distinct reply ids = %d, want %d", len(ids), calls)
}
for id := 1; id <= calls; id++ {
if !ids[float64(id)] {
t.Errorf("missing reply id %d", id)
}
}
}
type firstFailureWriter struct {
calls int
first error
later error
}
func (w *firstFailureWriter) Write([]byte) (int, error) {
w.calls++
if w.calls == 1 {
return 0, w.first
}
return 0, w.later
}
func TestConcurrentRepliesReturnTheFirstWriteFailure(t *testing.T) {
first := errors.New("stdout closed")
later := errors.New("a later failure")
out := &firstFailureWriter{first: first, later: later}
server := serverFor(&fakeCluster{}, Options{})
var in strings.Builder
for id := 1; id <= 16; id++ {
_, _ = fmt.Fprintf(&in, `{"jsonrpc":"2.0","id":%d,"method":"ping"}`+"\n", id)
}
err := server.Serve(context.Background(), strings.NewReader(in.String()), out)
if !errors.Is(err, first) {
t.Fatalf("serve error = %v, want the first write failure", err)
}
if out.calls != 1 {
t.Fatalf("writer calls = %d, want replies suppressed after the failure", out.calls)
}
}
type lockedBuffer struct {
mu sync.Mutex
body bytes.Buffer
written chan struct{}
}
func (l *lockedBuffer) Write(p []byte) (int, error) {
l.mu.Lock()
written, err := l.body.Write(p)
l.mu.Unlock()
if l.written != nil {
select {
case l.written <- struct{}{}:
default:
}
}
return written, err
}
func (l *lockedBuffer) read() string {
l.mu.Lock()
defer l.mu.Unlock()
return l.body.String()
}