From 50eddac5a1e3b7a6d50806d762d11a87ba22c1b5 Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Fri, 12 Sep 2025 14:42:07 +0000
Subject: [PATCH] mcp: speed up TestCommandTransportDuration

At 7s, this test was 70% of our test execution time on my computer.
Speed it up by mutating the defaultTerminateDuration for the purposes of
the test.

Also update our test workflow to use go1.25, now that it's out.
---
 .github/workflows/test.yml |  2 +-
 mcp/cmd.go                 |  4 +---
 mcp/cmd_export_test.go     | 20 ++++++++++++++++++++
 mcp/cmd_test.go            | 24 ++++++++++++++++++------
 4 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 mcp/cmd_export_test.go

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d8e2d31e..9ad50c1d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -39,7 +39,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: ["1.23", "1.24", "1.25.0-rc.3"]
+        go: ["1.23", "1.24", "1.25"]
     steps:
       - name: Check out code
         uses: actions/checkout@v4
diff --git a/mcp/cmd.go b/mcp/cmd.go
index 55e5cca6..b531eaf1 100644
--- a/mcp/cmd.go
+++ b/mcp/cmd.go
@@ -13,9 +13,7 @@ import (
 	"time"
 )
 
-const (
-	defaultTerminateDuration = 5 * time.Second
-)
+var defaultTerminateDuration = 5 * time.Second // mutable for testing
 
 // A CommandTransport is a [Transport] that runs a command and communicates
 // with it over stdin/stdout, using newline-delimited JSON.
diff --git a/mcp/cmd_export_test.go b/mcp/cmd_export_test.go
new file mode 100644
index 00000000..331e8bd0
--- /dev/null
+++ b/mcp/cmd_export_test.go
@@ -0,0 +1,20 @@
+// 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 "time"
+
+// This file exports some helpers for mutating internals of the command
+// transport for testing.
+
+// SetDefaultTerminateDuration sets the default command terminate duration,
+// and returns a function to reset it to the default.
+func SetDefaultTerminateDuration(d time.Duration) (reset func()) {
+	initial := defaultTerminateDuration
+	defaultTerminateDuration = d
+	return func() {
+		defaultTerminateDuration = initial
+	}
+}
diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go
index 146cbe1f..0df45708 100644
--- a/mcp/cmd_test.go
+++ b/mcp/cmd_test.go
@@ -257,25 +257,35 @@ func TestCommandTransportTerminateDuration(t *testing.T) {
 	}
 	requireExec(t)
 
+	// Unfortunately, since it does I/O, this test needs to rely on timing (we
+	// can't use synctest). However, we can still decreate the default
+	// termination duration to speed up the test.
+	const defaultDur = 50 * time.Millisecond
+	defer mcp.SetDefaultTerminateDuration(defaultDur)()
+
 	tests := []struct {
 		name            string
 		duration        time.Duration
+		wantMinDuration time.Duration
 		wantMaxDuration time.Duration
 	}{
 		{
 			name:            "default duration (zero)",
 			duration:        0,
-			wantMaxDuration: 6 * time.Second, // default 5s + buffer
+			wantMinDuration: defaultDur,
+			wantMaxDuration: 1 * time.Second, // default + buffer
 		},
 		{
 			name:            "below minimum duration",
-			duration:        500 * time.Millisecond,
-			wantMaxDuration: 6 * time.Second, // should use default 5s + buffer
+			duration:        -500 * time.Millisecond,
+			wantMinDuration: defaultDur,
+			wantMaxDuration: 1 * time.Second, // should use default + buffer
 		},
 		{
 			name:            "custom valid duration",
-			duration:        2 * time.Second,
-			wantMaxDuration: 3 * time.Second, // custom 2s + buffer
+			duration:        200 * time.Millisecond,
+			wantMinDuration: 200 * time.Millisecond,
+			wantMaxDuration: 1 * time.Second, // custom + buffer
 		},
 	}
 
@@ -306,7 +316,9 @@ func TestCommandTransportTerminateDuration(t *testing.T) {
 					t.Fatalf("Close() failed with unexpected error: %v", err)
 				}
 			}
-
+			if elapsed < tt.wantMinDuration {
+				t.Errorf("Close() took %v, expected at least %v", elapsed, tt.wantMinDuration)
+			}
 			if elapsed > tt.wantMaxDuration {
 				t.Errorf("Close() took %v, expected at most %v", elapsed, tt.wantMaxDuration)
 			}