1
1
package io.modelcontextprotocol.kotlin.sdk.integration
2
2
3
3
import io.ktor.client.HttpClient
4
- import io.ktor.server.application.ApplicationStopped
5
4
import io.ktor.server.application.install
6
5
import io.ktor.server.cio.CIOApplicationEngine
7
6
import io.ktor.server.engine.EmbeddedServer
@@ -20,12 +19,12 @@ import io.modelcontextprotocol.kotlin.sdk.client.mcpWebSocketTransport
20
19
import io.modelcontextprotocol.kotlin.sdk.server.Server
21
20
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
22
21
import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocket
22
+ import kotlin.test.Test
23
+ import kotlin.test.assertTrue
24
+ import kotlin.time.Duration.Companion.seconds
23
25
import kotlinx.coroutines.Dispatchers
24
26
import kotlinx.coroutines.test.runTest
25
27
import kotlinx.coroutines.withContext
26
- import kotlinx.coroutines.withTimeout
27
- import kotlin.test.Test
28
- import kotlin.test.assertTrue
29
28
import io.ktor.client.engine.cio.CIO as ClientCIO
30
29
import io.ktor.client.plugins.websocket.WebSockets as ClientWebSocket
31
30
import io.ktor.server.cio.CIO as ServerCIO
@@ -34,21 +33,19 @@ import io.ktor.server.websocket.WebSockets as ServerWebSockets
34
33
class WebSocketIntegrationTest {
35
34
36
35
@Test
37
- fun `client should be able to connect to websocket server 2` () = runTest {
36
+ fun `client should be able to connect to websocket server 2` () = runTest(timeout = 5 .seconds) {
38
37
var server: EmbeddedServer <CIOApplicationEngine , CIOApplicationEngine .Configuration >? = null
39
38
var client: Client ? = null
40
39
41
40
try {
42
41
withContext(Dispatchers .Default ) {
43
- withTimeout(1000 ) {
44
- server = initServer()
45
- val port = server.engine.resolvedConnectors().first().port
46
- client = initClient(serverPort = port)
47
- }
42
+ server = initServer()
43
+ val port = server.engine.resolvedConnectors().first().port
44
+ client = initClient(serverPort = port)
48
45
}
49
46
} finally {
50
47
client?.close()
51
- server?.stop (1000 , 2000 )
48
+ server?.stopSuspend (1000 , 2000 )
52
49
}
53
50
}
54
51
@@ -60,24 +57,22 @@ class WebSocketIntegrationTest {
60
57
* 3. Observe that Client A receives a response related to it.
61
58
*/
62
59
@Test
63
- fun `single websocket connection` () = runTest {
60
+ fun `single websocket connection` () = runTest(timeout = 5 .seconds) {
64
61
var server: EmbeddedServer <CIOApplicationEngine , CIOApplicationEngine .Configuration >? = null
65
62
var client: Client ? = null
66
63
67
64
try {
68
65
withContext(Dispatchers .Default ) {
69
- withTimeout(1000 ) {
70
- server = initServer()
71
- val port = server.engine.resolvedConnectors().first().port
72
- client = initClient(" Client A" , port)
73
-
74
- val promptA = getPrompt(client, " Client A" )
75
- assertTrue { " Client A" in promptA }
76
- }
66
+ server = initServer()
67
+ val port = server.engine.resolvedConnectors().first().port
68
+ client = initClient(" Client A" , port)
69
+
70
+ val promptA = getPrompt(client, " Client A" )
71
+ assertTrue { " Client A" in promptA }
77
72
}
78
73
} finally {
79
74
client?.close()
80
- server?.stop (1000 , 2000 )
75
+ server?.stopSuspend (1000 , 2000 )
81
76
}
82
77
}
83
78
@@ -90,32 +85,30 @@ class WebSocketIntegrationTest {
90
85
* 4. Observe that Client B (connection #2) receives a response related to sessionId#1.
91
86
*/
92
87
@Test
93
- fun `multiple websocket connections` () = runTest {
88
+ fun `multiple websocket connections` () = runTest(timeout = 5 .seconds) {
94
89
var server: EmbeddedServer <CIOApplicationEngine , CIOApplicationEngine .Configuration >? = null
95
90
var clientA: Client ? = null
96
91
var clientB: Client ? = null
97
92
98
93
try {
99
94
withContext(Dispatchers .Default ) {
100
- withTimeout(1000 ) {
101
- server = initServer()
102
- val port = server.engine.resolvedConnectors().first().port
103
- clientA = initClient(" Client A" , port)
104
- clientB = initClient(" Client B" ,port)
105
-
106
- // Step 3: Send a prompt request from Client A
107
- val promptA = getPrompt(clientA, " Client A" )
108
- // Step 4: Send a prompt request from Client B
109
- val promptB = getPrompt(clientB, " Client B" )
110
-
111
- assertTrue { " Client A" in promptA }
112
- assertTrue { " Client B" in promptB }
113
- }
95
+ server = initServer()
96
+ val port = server.engine.resolvedConnectors().first().port
97
+ clientA = initClient(" Client A" , port)
98
+ clientB = initClient(" Client B" , port)
99
+
100
+ // Step 3: Send a prompt request from Client A
101
+ val promptA = getPrompt(clientA, " Client A" )
102
+ // Step 4: Send a prompt request from Client B
103
+ val promptB = getPrompt(clientB, " Client B" )
104
+
105
+ assertTrue { " Client A" in promptA }
106
+ assertTrue { " Client B" in promptB }
114
107
}
115
108
} finally {
116
109
clientA?.close()
117
110
clientB?.close()
118
- server?.stop (1000 , 2000 )
111
+ server?.stopSuspend (1000 , 2000 )
119
112
}
120
113
}
121
114
@@ -178,10 +171,6 @@ class WebSocketIntegrationTest {
178
171
}
179
172
}
180
173
181
- ktorServer.monitor.subscribe(ApplicationStopped ) {
182
- println (" SD -- [T] ktor server has been stopped" )
183
- }
184
-
185
174
return ktorServer.startSuspend(wait = false )
186
175
}
187
176
0 commit comments