Skip to content

Commit e6b9411

Browse files
tiginamariatiginamaria
committed
Use awaitility for testing notifications and remove unnecessary server.close() calls
1 parent ce1d215 commit e6b9411

File tree

6 files changed

+89
-109
lines changed

6 files changed

+89
-109
lines changed

‎kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureNotificationService.kt‎

Lines changed: 31 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification
66
import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification
77
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification
88
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotificationParams
9+
import io.modelcontextprotocol.kotlin.sdk.types.ServerNotification
910
import io.modelcontextprotocol.kotlin.sdk.types.ToolListChangedNotification
1011
import kotlinx.atomicfu.atomic
1112
import kotlinx.atomicfu.getAndUpdate
@@ -18,6 +19,7 @@ import kotlinx.coroutines.Job
1819
import kotlinx.coroutines.SupervisorJob
1920
import kotlinx.coroutines.cancel
2021
import kotlinx.coroutines.channels.BufferOverflow
22+
import kotlinx.coroutines.channels.Channel
2123
import kotlinx.coroutines.flow.MutableSharedFlow
2224
import kotlinx.coroutines.flow.SharedFlow
2325
import kotlinx.coroutines.flow.takeWhile
@@ -58,6 +60,17 @@ private class SessionNotificationJob {
5860
private val resourceSubscriptions = atomic(persistentMapOf<FeatureKey, Long>())
5961
private val logger = KotlinLogging.logger {}
6062

63+
/**
64+
* Constructor for the SessionNotificationJob, responsible for processing notification events
65+
* and dispatching appropriate notifications to the provided server session. The job operates
66+
* within the given coroutine scope and begins handling events starting from the specified
67+
* timestamp.
68+
*
69+
* @param session The server session where notifications will be dispatched.
70+
* @param scope The coroutine scope in which this job operates.
71+
* @param events A shared flow of notification events that the job listens to.
72+
* @param fromTimestamp The timestamp from which the job starts processing events.
73+
*/
6174
constructor(
6275
session: ServerSession,
6376
scope: CoroutineScope,
@@ -71,23 +84,12 @@ private class SessionNotificationJob {
7184
is SendEvent -> {
7285
if (event.timestamp >= fromTimestamp) {
7386
when (val notification = event.notification) {
74-
is PromptListChangedNotification -> {
75-
logger.info {
76-
"Sending prompt list changed notification for sessionId: ${session.sessionId}"
77-
}
78-
session.notification(notification)
79-
}
80-
81-
is ResourceListChangedNotification -> {
87+
is PromptListChangedNotification,
88+
is ResourceListChangedNotification,
89+
is ToolListChangedNotification,
90+
-> {
8291
logger.info {
83-
"Sending resourse list changed notification for sessionId: ${session.sessionId}"
84-
}
85-
session.notification(notification)
86-
}
87-
88-
is ToolListChangedNotification -> {
89-
logger.info {
90-
"Sending tool list changed notification for sessionId: ${session.sessionId}"
92+
"Sending list changed notification for sessionId: ${session.sessionId}"
9193
}
9294
session.notification(notification)
9395
}
@@ -186,6 +188,7 @@ private class SessionNotificationJob {
186188
* - Resource updates pertaining to specific resources.
187189
*/
188190
internal class FeatureNotificationService(
191+
notificationBufferCapacity: Int = Channel.UNLIMITED,
189192
@OptIn(ExperimentalTime::class)
190193
private val clock: Clock = Clock.System,
191194
) {
@@ -197,8 +200,7 @@ internal class FeatureNotificationService(
197200

198201
/** Shared flow used to emit events within the feature notification service. */
199202
private val notificationEvents = MutableSharedFlow<NotificationEvent>(
200-
extraBufferCapacity = 100,
201-
replay = 0,
203+
extraBufferCapacity = notificationBufferCapacity,
202204
onBufferOverflow = BufferOverflow.SUSPEND,
203205
)
204206

@@ -210,57 +212,27 @@ internal class FeatureNotificationService(
210212

211213
private val logger = KotlinLogging.logger {}
212214

213-
/** Listener for tool feature events. */
214-
private val toolListChangedListener: FeatureListener by lazy {
215+
private fun featureListener(notificationProvider: (FeatureKey) -> ServerNotification): FeatureListener =
215216
object : FeatureListener {
216217
override fun onFeatureUpdated(featureKey: FeatureKey) {
217-
logger.info { "Emitting tool list changed notification" }
218-
emit(ToolListChangedNotification())
218+
val notification = notificationProvider(featureKey)
219+
logger.info { "Emitting notification: ${notification.method.value}" }
220+
emit(notification)
219221
}
220222
}
221-
}
223+
224+
/** Listener for tool feature events. */
225+
internal val toolListChangedListener: FeatureListener = featureListener { ToolListChangedNotification() }
222226

223227
/** Listener for prompt feature events. */
224-
private val promptListChangeListener: FeatureListener by lazy {
225-
object : FeatureListener {
226-
override fun onFeatureUpdated(featureKey: FeatureKey) {
227-
logger.info { "Emitting prompt list changed notification" }
228-
emit(PromptListChangedNotification())
229-
}
230-
}
231-
}
228+
internal val promptListChangedListener: FeatureListener = featureListener { PromptListChangedNotification() }
232229

233230
/** Listener for resource feature events. */
234-
private val resourceListChangedListener: FeatureListener by lazy {
235-
object : FeatureListener {
236-
override fun onFeatureUpdated(featureKey: FeatureKey) {
237-
logger.info { "Emitting resource list changed notification" }
238-
emit(ResourceListChangedNotification())
239-
}
240-
}
241-
}
242-
243-
/** Listener for resource update events. */
244-
private val resourceUpdatedListener: FeatureListener by lazy {
245-
object : FeatureListener {
246-
override fun onFeatureUpdated(featureKey: FeatureKey) {
247-
logger.info { "Emitting resource updated notification for feature key: $featureKey" }
248-
emit(ResourceUpdatedNotification(ResourceUpdatedNotificationParams(uri = featureKey)))
249-
}
250-
}
251-
}
252-
253-
/** Listener for the tool list changed events. */
254-
internal fun getToolListChangedListener(): FeatureListener = toolListChangedListener
255-
256-
/** Listener for the prompt list changed events. */
257-
internal fun getPromptListChangedListener(): FeatureListener = promptListChangeListener
258-
259-
/** Listener for the resource list changed events. */
260-
internal fun getResourceListChangedListener(): FeatureListener = resourceListChangedListener
231+
internal val resourceListChangedListener: FeatureListener = featureListener { ResourceListChangedNotification() }
261232

262233
/** Listener for resource update events. */
263-
internal fun getResourceUpdateListener(): FeatureListener = resourceUpdatedListener
234+
internal val resourceUpdatedListener: FeatureListener =
235+
featureListener { ResourceUpdatedNotification(ResourceUpdatedNotificationParams(uri = it)) }
264236

265237
/**
266238
* Subscribes session to list changed notifications for all features and resource update notifications.

‎kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,20 +113,20 @@ public open class Server(
113113

114114
private val toolRegistry = FeatureRegistry<RegisteredTool>("Tool").apply {
115115
if (options.capabilities.tools?.listChanged ?: false) {
116-
addListener(notificationService.getToolListChangedListener())
116+
addListener(notificationService.toolListChangedListener)
117117
}
118118
}
119119
private val promptRegistry = FeatureRegistry<RegisteredPrompt>("Prompt").apply {
120120
if (options.capabilities.prompts?.listChanged ?: false) {
121-
addListener(notificationService.getPromptListChangedListener())
121+
addListener(notificationService.promptListChangedListener)
122122
}
123123
}
124124
private val resourceRegistry = FeatureRegistry<RegisteredResource>("Resource").apply {
125125
if (options.capabilities.resources?.listChanged ?: false) {
126-
addListener(notificationService.getResourceListChangedListener())
126+
addListener(notificationService.resourceListChangedListener)
127127
}
128128
if (options.capabilities.resources?.subscribe ?: false) {
129-
addListener(notificationService.getResourceUpdateListener())
129+
addListener(notificationService.resourceUpdatedListener)
130130
}
131131
}
132132

‎kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsNotificationTest.kt‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification
77
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
88
import kotlinx.coroutines.CompletableDeferred
99
import kotlinx.coroutines.test.runTest
10+
import org.awaitility.kotlin.await
11+
import org.awaitility.kotlin.untilAsserted
1012
import org.junit.jupiter.api.Test
1113
import kotlin.test.assertEquals
1214
import kotlin.test.assertFalse
@@ -38,14 +40,14 @@ class ServerPromptsNotificationTest : AbstractServerFeaturesTest() {
3840

3941
// Remove the prompt
4042
val result = server.removePrompt(testPrompt.name)
41-
// Close the server to stop processing further events and flush notifications
42-
server.close()
4343

4444
// Verify the prompt was removed
4545
assertTrue(result, "Prompt should be removed successfully")
4646

4747
// Verify that the notification was sent
48-
assertTrue(promptListChangedNotificationReceived, "Notification should be sent when prompt is added")
48+
await untilAsserted {
49+
assertTrue(promptListChangedNotificationReceived, "Notification should be sent when prompt is added")
50+
}
4951
}
5052

5153
@Test
@@ -76,18 +78,18 @@ class ServerPromptsNotificationTest : AbstractServerFeaturesTest() {
7678

7779
// Remove the prompts
7880
val result = server.removePrompts(listOf(testPrompt1.name, testPrompt2.name))
79-
// Close the server to stop processing further events and flush notifications
80-
server.close()
8181

8282
// Verify the prompts were removed
8383
assertEquals(2, result, "Both prompts should be removed")
8484

8585
// Verify that the notifications were sent twice
86-
assertEquals(
87-
4,
88-
promptListChangedNotificationReceivedCount,
89-
"Two notifications should be sent when prompts are added and two when removed",
90-
)
86+
await untilAsserted {
87+
assertEquals(
88+
4,
89+
promptListChangedNotificationReceivedCount,
90+
"Two notifications should be sent when prompts are added and two when removed",
91+
)
92+
}
9193
}
9294

9395
@Test
@@ -101,11 +103,14 @@ class ServerPromptsNotificationTest : AbstractServerFeaturesTest() {
101103

102104
// Try to remove a non-existent prompt
103105
val result = server.removePrompt("non-existent-prompt")
104-
// Close the server to stop processing further events and flush notifications
105-
server.close()
106106

107107
// Verify the result
108108
assertFalse(result, "Removing non-existent prompt should return false")
109-
assertFalse(promptListChangedNotificationReceived, "No notification should be sent when prompt doesn't exist")
109+
await untilAsserted {
110+
assertFalse(
111+
promptListChangedNotificationReceived,
112+
"No notification should be sent when prompt doesn't exist",
113+
)
114+
}
110115
}
111116
}

‎kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationSubscribeTest.kt‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import io.modelcontextprotocol.kotlin.sdk.types.SubscribeRequestParams
1010
import io.modelcontextprotocol.kotlin.sdk.types.TextResourceContents
1111
import kotlinx.coroutines.CompletableDeferred
1212
import kotlinx.coroutines.test.runTest
13+
import org.awaitility.kotlin.await
14+
import org.awaitility.kotlin.untilAsserted
1315
import org.junit.jupiter.api.Test
1416
import kotlin.test.assertEquals
1517
import kotlin.test.assertFalse
@@ -26,7 +28,6 @@ class ServerResourcesNotificationSubscribeTest : AbstractServerFeaturesTest() {
2628
val notifications = mutableListOf<ResourceUpdatedNotification>()
2729
client.setNotificationHandler<ResourceUpdatedNotification>(Method.Defined.NotificationsResourcesUpdated) {
2830
notifications.add(it)
29-
println(it)
3031
CompletableDeferred(Unit)
3132
}
3233

@@ -72,15 +73,14 @@ class ServerResourcesNotificationSubscribeTest : AbstractServerFeaturesTest() {
7273

7374
// Remove the resource
7475
val result = server.removeResource(testResourceUri1)
75-
// Close the server to stop processing further events and flush notifications
76-
server.close()
7776

7877
// Verify the resource was removed
7978
assertTrue(result, "Resource should be removed successfully")
8079

81-
println(notifications.map { it.params.uri })
8280
// Verify that the notification was sent
83-
assertEquals(1, notifications.size, "Notification should be sent when resource 1 was deleted")
81+
await untilAsserted {
82+
assertEquals(1, notifications.size, "Notification should be sent when resource 1 was deleted")
83+
}
8484
assertEquals(testResourceUri1, notifications[0].params.uri, "Notification should contain the resource 1 URI")
8585
}
8686

@@ -136,16 +136,20 @@ class ServerResourcesNotificationSubscribeTest : AbstractServerFeaturesTest() {
136136
// Remove the resource
137137
val result1 = server.removeResource(testResourceUri1)
138138
val result2 = server.removeResource(testResourceUri2)
139-
// Close the server to stop processing further events and flush notifications
140-
server.close()
141139

142140
// Verify the resource was removed
143141
assertTrue(result1, "Resource 1 should be removed successfully")
144142
assertTrue(result2, "Resource 2 should be removed successfully")
145143

146144
println(notifications.map { it.params.uri })
147145
// Verify that the notification was sent
148-
assertEquals(2, notifications.size, "Notification should be sent when resource 1 and resource 2 was deleted")
146+
await untilAsserted {
147+
assertEquals(
148+
2,
149+
notifications.size,
150+
"Notification should be sent when resource 1 and resource 2 was deleted",
151+
)
152+
}
149153

150154
val deletedResources = listOf(notifications[0].params.uri, notifications[1].params.uri)
151155
assertTrue(
@@ -171,8 +175,6 @@ class ServerResourcesNotificationSubscribeTest : AbstractServerFeaturesTest() {
171175

172176
// Try to remove a non-existent resource
173177
val result = server.removeResource("non-existent-resource")
174-
// Close the server to stop processing further events and flush notifications
175-
server.close()
176178

177179
// Verify the result
178180
assertFalse(result, "Removing non-existent resource should return false")

‎kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationTest.kt‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import org.junit.jupiter.api.Test
1111
import kotlin.test.assertEquals
1212
import kotlin.test.assertFalse
1313
import kotlin.test.assertTrue
14+
import org.awaitility.kotlin.await
15+
import org.awaitility.kotlin.untilAsserted
1416

1517
class ServerResourcesNotificationTest : AbstractServerFeaturesTest() {
1618

@@ -50,14 +52,14 @@ class ServerResourcesNotificationTest : AbstractServerFeaturesTest() {
5052

5153
// Remove the resource
5254
val result = server.removeResource(testResourceUri)
53-
// Close the server to stop processing further events and flush notifications
54-
server.close()
5555

5656
// Verify the resource was removed
5757
assertTrue(result, "Resource should be removed successfully")
5858

5959
// Verify that the notification was sent
60-
assertTrue(resourceListChangedNotificationReceived, "Notification should be sent when resource is added")
60+
await untilAsserted {
61+
assertTrue(resourceListChangedNotificationReceived, "Notification should be sent when resource is added")
62+
}
6163
}
6264

6365
@Test
@@ -110,18 +112,18 @@ class ServerResourcesNotificationTest : AbstractServerFeaturesTest() {
110112

111113
// Remove the resources
112114
val result = server.removeResources(listOf(testResourceUri1, testResourceUri2))
113-
// Close the server to stop processing further events and flush notifications
114-
server.close()
115115

116116
// Verify the resources were removed
117117
assertEquals(2, result, "Both resources should be removed")
118118

119119
// Verify that the notifications were sent twice
120-
assertEquals(
121-
4,
122-
resourceListChangedNotificationReceivedCount,
123-
"Two notifications should be sent when resources are added and two when removed",
124-
)
120+
await untilAsserted {
121+
assertEquals(
122+
4,
123+
resourceListChangedNotificationReceivedCount,
124+
"Two notifications should be sent when resources are added and two when removed",
125+
)
126+
}
125127
}
126128

127129
@Test
@@ -137,8 +139,6 @@ class ServerResourcesNotificationTest : AbstractServerFeaturesTest() {
137139

138140
// Try to remove a non-existent resource
139141
val result = server.removeResource("non-existent-resource")
140-
// Close the server to stop processing further events and flush notifications
141-
server.close()
142142

143143
// Verify the result
144144
assertFalse(result, "Removing non-existent resource should return false")

0 commit comments

Comments
 (0)