Skip to content

Commit bc28e6c

Browse files
kpavlovkpavlov
authored
Refactor streaming http tests
To provide better DSL for integration tests
1 parent d60e8b6 commit bc28e6c

File tree

4 files changed

+233
-105
lines changed

4 files changed

+233
-105
lines changed

‎kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/AbstractStreamableHttpClientTest.kt‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package io.modelcontextprotocol.kotlin.sdk.client
2+
3+
import io.ktor.client.HttpClient
4+
import io.ktor.client.engine.apache5.Apache5
5+
import io.ktor.client.plugins.logging.LogLevel
6+
import io.ktor.client.plugins.logging.Logging
7+
import io.ktor.client.plugins.sse.SSE
8+
import org.junit.jupiter.api.AfterEach
9+
import org.junit.jupiter.api.TestInstance
10+
11+
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
12+
internal abstract class AbstractStreamableHttpClientTest {
13+
14+
// start mokksy on random port
15+
protected val mockMcp: MockMcp = MockMcp(verbose = true)
16+
17+
@AfterEach
18+
fun afterEach() {
19+
mockMcp.checkForUnmatchedRequests()
20+
}
21+
22+
protected suspend fun connect(client: Client) {
23+
client.connect(
24+
StreamableHttpClientTransport(
25+
url = mockMcp.url,
26+
client = HttpClient(Apache5) {
27+
install(SSE)
28+
install(Logging) {
29+
level = LogLevel.ALL
30+
}
31+
},
32+
),
33+
)
34+
}
35+
}

‎kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/MockMcp.kt‎

Lines changed: 177 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,25 @@
11
package io.modelcontextprotocol.kotlin.sdk.client
22

3+
import dev.mokksy.mokksy.BuildingStep
34
import dev.mokksy.mokksy.Mokksy
45
import dev.mokksy.mokksy.StubConfiguration
56
import io.ktor.http.ContentType
67
import io.ktor.http.HttpMethod
78
import io.ktor.http.HttpStatusCode
89
import io.ktor.sse.ServerSentEvent
910
import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
11+
import io.modelcontextprotocol.kotlin.sdk.RequestId
1012
import kotlinx.coroutines.flow.Flow
13+
import kotlinx.serialization.json.Json
14+
import kotlinx.serialization.json.JsonObject
15+
import kotlinx.serialization.json.JsonPrimitive
16+
import kotlinx.serialization.json.buildJsonObject
17+
import kotlinx.serialization.json.contentOrNull
18+
import kotlinx.serialization.json.jsonObject
19+
import kotlinx.serialization.json.jsonPrimitive
20+
import kotlinx.serialization.json.putJsonObject
21+
22+
const val MCP_SESSION_ID_HEADER = "Mcp-Session-Id"
1123

1224
/**
1325
* High-level helper for simulating an MCP server over Streaming HTTP transport with Server-Sent Events (SSE),
@@ -26,51 +38,188 @@ internal class MockMcp(verbose: Boolean = false) {
2638
mokksy.checkForUnmatchedRequests()
2739
}
2840

29-
val url = mokksy.baseUrl() + "/mcp"
41+
val url = "${mokksy.baseUrl()}/mcp"
3042

3143
@Suppress("LongParameterList")
44+
fun onInitialize(
45+
clientName: String? = null,
46+
sessionId: String,
47+
protocolVersion: String = "2025-03-26",
48+
serverName: String = "Mock MCP Server",
49+
serverVersion: String = "1.0.0",
50+
capabilities: JsonObject = buildJsonObject {
51+
putJsonObject("tools") {
52+
put("listChanged", JsonPrimitive(false))
53+
}
54+
},
55+
) {
56+
val predicates = if (clientName != null) {
57+
arrayOf<(JSONRPCRequest?) -> Boolean>({
58+
it?.params?.jsonObject
59+
?.get("clientInfo")?.jsonObject
60+
?.get("name")?.jsonPrimitive
61+
?.contentOrNull == clientName
62+
})
63+
} else {
64+
emptyArray()
65+
}
66+
67+
handleWithResult(
68+
jsonRpcMethod = "initialize",
69+
sessionId = sessionId,
70+
bodyPredicates = predicates,
71+
// language=json
72+
result = """
73+
{
74+
"capabilities": $capabilities,
75+
"protocolVersion": "$protocolVersion",
76+
"serverInfo": {
77+
"name": "$serverName",
78+
"version": "$serverVersion"
79+
},
80+
"_meta": {
81+
"foo": "bar"
82+
}
83+
}
84+
""".trimIndent(),
85+
)
86+
}
87+
3288
fun onJSONRPCRequest(
89+
httpMethod: HttpMethod = HttpMethod.Post,
90+
jsonRpcMethod: String,
91+
expectedSessionId: String? = null,
92+
vararg bodyPredicates: (JSONRPCRequest) -> Boolean,
93+
): BuildingStep<JSONRPCRequest> = mokksy.method(
94+
configuration = StubConfiguration(removeAfterMatch = true),
95+
httpMethod = httpMethod,
96+
requestType = JSONRPCRequest::class,
97+
) {
98+
path("/mcp")
99+
expectedSessionId?.let {
100+
containsHeader(MCP_SESSION_ID_HEADER, it)
101+
}
102+
bodyMatchesPredicate(
103+
description = "JSON-RPC version is '2.0'",
104+
predicate =
105+
{
106+
it!!.jsonrpc == "2.0"
107+
},
108+
)
109+
bodyMatchesPredicate(
110+
description = "JSON-RPC Method should be '$jsonRpcMethod'",
111+
predicate =
112+
{
113+
it!!.method == jsonRpcMethod
114+
},
115+
)
116+
bodyPredicates.forEach { predicate ->
117+
bodyMatchesPredicate(predicate = { predicate.invoke(it!!) })
118+
}
119+
}
120+
121+
@Suppress("LongParameterList")
122+
fun handleWithResult(
33123
httpMethod: HttpMethod = HttpMethod.Post,
34124
jsonRpcMethod: String,
35125
expectedSessionId: String? = null,
36126
sessionId: String,
37127
contentType: ContentType = ContentType.Application.Json,
38128
statusCode: HttpStatusCode = HttpStatusCode.OK,
39-
bodyBuilder: () -> String,
129+
vararg bodyPredicates: (JSONRPCRequest) -> Boolean,
130+
result: () -> JsonObject,
40131
) {
41-
mokksy.method(
42-
configuration = StubConfiguration(removeAfterMatch = true),
132+
onJSONRPCRequest(
43133
httpMethod = httpMethod,
44-
requestType = JSONRPCRequest::class,
45-
) {
46-
path("/mcp")
47-
expectedSessionId?.let {
48-
containsHeader("Mcp-Session-Id", it)
134+
jsonRpcMethod = jsonRpcMethod,
135+
expectedSessionId = expectedSessionId,
136+
bodyPredicates = bodyPredicates,
137+
) respondsWith {
138+
val requestId = when (request.body.id) {
139+
is RequestId.NumberId -> (request.body.id as RequestId.NumberId).value.toString()
140+
is RequestId.StringId -> "\"${(request.body.id as RequestId.StringId).value}\""
49141
}
50-
bodyMatchesPredicates(
51-
{
52-
it!!.method == jsonRpcMethod
53-
},
54-
{
55-
it!!.jsonrpc == "2.0"
56-
},
57-
)
58-
} respondsWith {
142+
val resultObject = result!!.invoke()
143+
// language=json
144+
body = """
145+
{
146+
"jsonrpc": "2.0",
147+
"id": $requestId,
148+
"result": $resultObject
149+
}
150+
""".trimIndent()
151+
this.contentType = contentType
152+
headers += MCP_SESSION_ID_HEADER to sessionId
153+
httpStatus = statusCode
154+
}
155+
}
156+
157+
@Suppress("LongParameterList")
158+
fun handleWithResult(
159+
httpMethod: HttpMethod = HttpMethod.Post,
160+
jsonRpcMethod: String,
161+
expectedSessionId: String? = null,
162+
sessionId: String,
163+
contentType: ContentType = ContentType.Application.Json,
164+
statusCode: HttpStatusCode = HttpStatusCode.OK,
165+
vararg bodyPredicates: (JSONRPCRequest) -> Boolean,
166+
result: String,
167+
) {
168+
handleWithResult(
169+
httpMethod = httpMethod,
170+
jsonRpcMethod = jsonRpcMethod,
171+
expectedSessionId = expectedSessionId,
172+
sessionId = sessionId,
173+
contentType = contentType,
174+
statusCode = statusCode,
175+
bodyPredicates = bodyPredicates,
176+
result = {
177+
Json.parseToJsonElement(result).jsonObject
178+
},
179+
)
180+
}
181+
182+
@Suppress("LongParameterList")
183+
fun handleJSONRPCRequest(
184+
httpMethod: HttpMethod = HttpMethod.Post,
185+
jsonRpcMethod: String,
186+
expectedSessionId: String? = null,
187+
sessionId: String,
188+
contentType: ContentType = ContentType.Application.Json,
189+
statusCode: HttpStatusCode = HttpStatusCode.OK,
190+
vararg bodyPredicates: (JSONRPCRequest?) -> Boolean,
191+
bodyBuilder: () -> String = { "" },
192+
) {
193+
onJSONRPCRequest(
194+
httpMethod = httpMethod,
195+
jsonRpcMethod = jsonRpcMethod,
196+
expectedSessionId = expectedSessionId,
197+
bodyPredicates = bodyPredicates,
198+
) respondsWith {
59199
body = bodyBuilder.invoke()
60200
this.contentType = contentType
61-
headers += "Mcp-Session-Id" to sessionId
201+
headers += MCP_SESSION_ID_HEADER to sessionId
62202
httpStatus = statusCode
63203
}
64204
}
65205

66-
fun onSubscribeWithGet(sessionId: String, block: () -> Flow<ServerSentEvent>) {
67-
mokksy.get(name = "MCP GETs", requestType = Any::class) {
68-
path("/mcp")
69-
containsHeader("Mcp-Session-Id", sessionId)
70-
containsHeader("Accept", "application/json,text/event-stream")
71-
containsHeader("Cache-Control", "no-store")
72-
} respondsWithSseStream {
73-
headers += "Mcp-Session-Id" to sessionId
206+
fun onSubscribe(httpMethod: HttpMethod = HttpMethod.Post, sessionId: String): BuildingStep<Any> = mokksy.method(
207+
httpMethod = httpMethod,
208+
name = "MCP GETs",
209+
requestType = Any::class,
210+
) {
211+
path("/mcp")
212+
containsHeader(MCP_SESSION_ID_HEADER, sessionId)
213+
containsHeader("Accept", "application/json,text/event-stream")
214+
containsHeader("Cache-Control", "no-store")
215+
}
216+
217+
fun handleSubscribeWithGet(sessionId: String, block: () -> Flow<ServerSentEvent>) {
218+
onSubscribe(
219+
httpMethod = HttpMethod.Get,
220+
sessionId = sessionId,
221+
) respondsWithSseStream {
222+
headers += MCP_SESSION_ID_HEADER to sessionId
74223
this.flow = block.invoke()
75224
}
76225
}
@@ -81,7 +230,7 @@ internal class MockMcp(verbose: Boolean = false) {
81230
requestType = JSONRPCRequest::class,
82231
) {
83232
path("/mcp")
84-
containsHeader("Mcp-Session-Id", sessionId)
233+
containsHeader(MCP_SESSION_ID_HEADER, sessionId)
85234
} respondsWith {
86235
body = null
87236
}

0 commit comments

Comments
 (0)