Skip to content

Commit 5d03db0

Browse files
committed
fix(test): Refactor flaky StdioClientTransport tests to use real stdin pipes
- Replace mock-based stdin with real input/output pipes for better parity with actual usage. - Ensure proper cleanup of resources (stdin and stdout pipes). - Improve test stability by migrating from `runTest` to `runBlocking` for real-time I/O operations.
1 parent adc4150 commit 5d03db0

File tree

1 file changed

+41
-24
lines changed

1 file changed

+41
-24
lines changed

kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/stdio/StdioClientTransportErrorHandlingTest.kt

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,27 +7,33 @@ import io.kotest.matchers.shouldBe
77
import io.kotest.matchers.types.shouldBeInstanceOf
88
import io.kotest.matchers.types.shouldBeSameInstanceAs
99
import io.mockk.coEvery
10-
import io.mockk.every
1110
import io.mockk.mockk
1211
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport
1312
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
1413
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
1514
import io.modelcontextprotocol.kotlin.sdk.types.McpException
1615
import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode
1716
import kotlinx.coroutines.CancellationException
17+
import kotlinx.coroutines.Dispatchers
1818
import kotlinx.coroutines.channels.Channel
1919
import kotlinx.coroutines.channels.ClosedSendChannelException
2020
import kotlinx.coroutines.delay
21+
import kotlinx.coroutines.runBlocking
2122
import kotlinx.coroutines.test.runTest
2223
import kotlinx.io.Buffer
24+
import kotlinx.io.asSource
25+
import kotlinx.io.buffered
2326
import kotlinx.io.writeString
2427
import org.junit.jupiter.params.ParameterizedTest
2528
import org.junit.jupiter.params.provider.Arguments
2629
import org.junit.jupiter.params.provider.MethodSource
30+
import java.io.PipedInputStream
31+
import java.io.PipedOutputStream
2732
import java.util.stream.Stream
2833
import kotlin.concurrent.atomics.AtomicBoolean
2934
import kotlin.concurrent.atomics.ExperimentalAtomicApi
3035
import kotlin.test.Test
36+
import kotlin.time.Duration.Companion.milliseconds
3137
import kotlin.time.Duration.Companion.seconds
3238

3339
/**
@@ -39,17 +45,23 @@ class StdioClientTransportErrorHandlingTest {
3945

4046
@OptIn(ExperimentalAtomicApi::class)
4147
@Test
42-
fun `should continue on stderr EOF`() = runTest {
43-
val stderrBuffer = Buffer()
48+
fun `should continue on stderr EOF`(): Unit = runBlocking(Dispatchers.IO) {
4449
// Empty stderr = immediate EOF
50+
val stderrBuffer = Buffer()
51+
52+
// Create a pipe for stdin that stays open (simulates real stdin behavior)
53+
val pipedOutputStream = PipedOutputStream()
54+
val pipedInputStream = PipedInputStream(pipedOutputStream)
55+
56+
// Write one message to stdin
57+
pipedOutputStream.write("""data: {"jsonrpc":"2.0","method":"ping","id":1}\n\n""".toByteArray())
58+
pipedOutputStream.flush()
59+
// Keep the pipe open by not closing pipedOutputStream - this prevents stdin EOF
4560

46-
val inputBuffer = createNonEmptyBuffer {
47-
"""data: {"jsonrpc":"2.0","method":"ping","id":1}\n\n"""
48-
}
4961
val outputBuffer = Buffer()
5062

5163
transport = StdioClientTransport(
52-
input = inputBuffer,
64+
input = pipedInputStream.asSource().buffered(),
5365
output = outputBuffer,
5466
error = stderrBuffer,
5567
)
@@ -59,15 +71,21 @@ class StdioClientTransportErrorHandlingTest {
5971

6072
transport.start()
6173

62-
// Stderr EOF should not close transport
63-
// Use eventually to handle timing differences across platforms (especially Windows)
74+
// Wait for stderr EOF and stdin message to be processed
75+
delay(500.milliseconds)
76+
77+
// Transport should still be alive because stdin is still open (not EOF'd)
78+
closeCalled.load() shouldBe false
79+
80+
// Close pipes to trigger stdin EOF.
81+
// `transport.close()` cann't help here, since the underlying Java read() is blocked on I/O operation
82+
pipedOutputStream.close()
83+
pipedInputStream.close()
84+
85+
// Transport should close when stdin EOF is detected
6486
eventually(2.seconds) {
65-
// Wait for stderr to be processed, then verify transport is still open
66-
closeCalled.load() shouldBe false
87+
closeCalled.load() shouldBe true
6788
}
68-
69-
transport.close()
70-
closeCalled.load() shouldBe true
7189
}
7290

7391
@Test
@@ -162,18 +180,13 @@ class StdioClientTransportErrorHandlingTest {
162180
fun `Send should handle exceptions`(throwable: Throwable, shouldWrap: Boolean, expectedCode: Int?) = runTest {
163181
val sendChannel: Channel<JSONRPCMessage> = mockk(relaxed = true)
164182

165-
// Create a stdin Source that never returns EOF to prevent transport from closing
166-
// Use mockk since Source is a sealed interface
167-
val stdin: kotlinx.io.Source = mockk(relaxed = true)
168-
every { stdin.readAtMostTo(any<Buffer>(), any()) } coAnswers {
169-
// Return 0 to indicate no data available (but not EOF)
170-
// This keeps the transport alive without blocking
171-
delay(10) // Small delay to prevent busy-waiting
172-
0L
173-
}
183+
// Create stdin pipe that stays open to prevent transport from closing
184+
val pipedOutputStream = PipedOutputStream()
185+
val pipedInputStream = PipedInputStream(pipedOutputStream)
186+
// Keep pipe open (don't write or close) - stdin will block on read, not EOF
174187

175188
transport = StdioClientTransport(
176-
input = stdin,
189+
input = pipedInputStream.asSource().buffered(),
177190
output = Buffer(),
178191
sendChannel = sendChannel,
179192
)
@@ -195,6 +208,10 @@ class StdioClientTransportErrorHandlingTest {
195208
} else {
196209
exception shouldBeSameInstanceAs throwable
197210
}
211+
212+
// Cleanup
213+
pipedOutputStream.close()
214+
pipedInputStream.close()
198215
}
199216

200217
fun createNonEmptyBuffer(block: () -> String): Buffer {

0 commit comments

Comments
 (0)