WebSocketClient frozen when connecting to WebSocket (Spring WebFlux) - spring

Question

I have written a sample to demo client/server communication with WebSocket protocol.
The server code:
#SpringBootApplication
class WebSocketServerApplication {
#Bean
fun webSocketMapping(mapper: ObjectMapper): HandlerMapping? {
val map = mapOf("/ws/messages" to ChatSocketHandler(mapper))
val simpleUrlHandlerMapping = SimpleUrlHandlerMapping().apply {
urlMap = map
order = 10
}
return simpleUrlHandlerMapping
}
#Bean
fun handlerAdapter(): WebSocketHandlerAdapter = WebSocketHandlerAdapter()
}
fun main(args: Array<String>) {
runApplication<WebSocketServerApplication>(*args)
}
class ChatSocketHandler(val mapper: ObjectMapper) : WebSocketHandler {
val sink = Sinks.replay<Message>(100);
val outputMessages: Flux<Message> = sink.asFlux();
override fun handle(session: WebSocketSession): Mono<Void> {
println("handling WebSocketSession...")
session.receive()
.map { it.payloadAsText }
.map { Message(id= UUID.randomUUID().toString(), body = it, sentAt = Instant.now()) }
.doOnNext { println(it) }
.subscribe(
{ message: Message -> sink.next(message) },
{ error: Throwable -> sink.error(error) }
);
return session.send(
Mono.delay(Duration.ofMillis(100))
.thenMany(outputMessages.map { session.textMessage(toJson(it)) })
)
}
fun toJson(message: Message): String = mapper.writeValueAsString(message)
}
data class Message #JsonCreator constructor(
#JsonProperty("id") var id: String? = null,
#JsonProperty("body") var body: String,
#JsonProperty("sentAt") var sentAt: Instant = Instant.now()
)
I have provided a client written in Angular, it works well, the codes is here.
When trying to a test for the server.
#SpringBootTest()
class WebsocketServerApplicationTests {
lateinit var client: WebSocketClient;
#Autowired
lateinit var mapper: ObjectMapper;
#BeforeEach
fun setup() {
this.client = ReactorNettyWebSocketClient()
}
#Test
fun contextLoads() {
val replay = Sinks.replay<Message>(10)
client.execute(
URI("ws://localhost:8080/ws/messages")
) { session: WebSocketSession ->
println("Starting to send messages")
session.receive()
.map { mapper.readValue(it.payloadAsText, Message::class.java) }
.subscribe { replay.next(it) }
session.send(
Mono.delay(Duration.ofSeconds(1)).thenMany(
Flux.just("test message", "test message2")
.map(session::textMessage)
)
).then()
}.subscribe()
StepVerifier.create(replay.asFlux().takeLast(2))
.consumeNextWith { it -> assertThat(it.body).isEqualTo("test message") }
.consumeNextWith { it -> assertThat(it.body).isEqualTo("test message2") }
.verifyComplete()
}
}
When starting up the application, run the test, it is frozen, not work as expected.

Answer 1

The problem is on the test side.
Complete your stream to make takeLast(n) working
First of all, you expect to take the last 2 elements from the stream. However, that is going to happen when and only when there is onComplete signal, which let the Flux.takeLast know that there is the end of the stream, so the last n observed elements are last.
In your code, you listen to the WebsocketInbound messages and send them to the ReplaySink. However, the FluxSink#complete message is never called, which means takeLast(2) will hang forever as expected.
Solution
On the one hand, the solution seems to be obvious:
session.receive()
.map { mapper.readValue(it.payloadAsText, Message::class.java) }
.subscribe ({ replay.next(it) }, { replay.error(it) }, { replay.complete() })
However, there is might be a trick:
.receive sends a terminal signal only when the WebSocket connection is closed.
Therefore, in order to receive a terminal signal, please ensure that the server closes the connection on its side. Otherwise, the test will still hang waiting for the final terminal signal.
If the connection close is not desired, try to simply use .take(2).

Answer 2

Finally fixed this issue myself after reading some posts on stackoverflow and source codes of testing reactive WebSocket in the spring framework.
Spring reactive ReactorNettyWebSocketClient not logging anything
How to use Spring Reactive WebSocket and transform it into the Flux stream?
WebSocketIntegrationTests
#SpringBootTest()
class WebSocketServerApplicationTests {
lateinit var client: WebSocketClient
#Autowired
lateinit var mapper: ObjectMapper
#BeforeEach
fun setup() {
this.client = ReactorNettyWebSocketClient()
}
#Test
fun contextLoads() {
val replay = Processors.replay<Message>(100)
try {
client.execute(
URI("ws://localhost:8080/ws/messages")
) { session: WebSocketSession ->
val receiveMono = session.receive()
.map { mapper.readValue(it.payloadAsText, Message::class.java) }
.log("received from server::")
.subscribeWith(replay)
.then()
session
.send(
Mono.delay(Duration.ofMillis(500)).thenMany(
Flux.just("test message", "test message2")
.map(session::textMessage)
)
)
.then(receiveMono)
}.block(Duration.ofSeconds(5L))
// assert
assertThat(replay.blockLast(Duration.ofSeconds(5L))?.body).isEqualTo("test message2")
} catch (e: Exception) {
println(e.message)
}
}
}