Java, Spring Boot에서 WebSocket 테스트하기

Spring Boot에서 WebSocket 테스트하기

Spring Boot 애플리케이션에서 WebSocket 기능을 구현한 후, 테스트를 진행했습니다. @SpringBootTest와 WebSocketStompClient를 활용하여 WebSocket 통신을 테스트하고, AssertJ를 사용하여 검증했습니다.

Spring Boot 테스트 설정 및 포트 구성

WebSocket 테스트를 위해 @SpringBootTest 어노테이션을 사용하여 통합 테스트 환경을 구성했습니다. 이때, webEnvironment 속성을 SpringBootTest.WebEnvironment.RANDOM_PORT로 설정했습니다. DEFINED_PORT를 사용하여 지정된 포트를 사용할 수도 있습니다.

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
public class WebSocketTest {
    @LocalServerPort
    private int port;

    // 테스트에 필요한 필드 및 메서드 정의
}

@LocalServerPort 어노테이션을 통해 할당된 포트를 주입받고 클라이언트 설정 시 사용했습니다

WebSocket 클라이언트 설정

테스트에서 WebSocket 클라이언트 역할을 수행하기 위해 WebSocketStompClient를 설정합니다. 이 클라이언트는 STOMP 프로토콜을 사용하여 서버와 통신합니다.

private WebSocketStompClient stompClient;

@BeforeEach
void setUp() {
    // SockJS 클라이언트 생성
    List<Transport> transports = List.of(new WebSocketTransport(new StandardWebSocketClient()));
    SockJsClient sockJsClient = new SockJsClient(transports);

    // WebSocketStompClient 생성 및 설정
    stompClient = new WebSocketStompClient(sockJsClient);
    stompClient.setMessageConverter(new MappingJackson2MessageConverter());
}

SockJsClient와 WebSocketTransport를 사용하여 다양한 전송 방식을 지원하며, MappingJackson2MessageConverter를 통해 JSON 형식의 메시지를 처리했습니다.

BlockingQueue와 CompletableFuture 비교

테스트에서 비동기적으로 수신되는 메시지를 처리하기 위해 BlockingQueue를 사용습니다. 이는 메시지를 대기열에 넣고, 테스트 코드에서 이를 가져와 검증하는 방식입니다.

private BlockingQueue<MessageDto> blockingQueue;

@BeforeEach
void setUp() {
    blockingQueue = new LinkedBlockingQueue<>();
    // stompClient 설정 코드
}

또한, BlockingQueue외에 CompletableFuture를 사용하여 비동기 작업의 완료를 기다리고 결과를 처리할 수 있습니다.

CompletableFuture<MessageDto> completableFuture = new CompletableFuture<>();

// 메시지 수신 시
completableFuture.complete(receivedMessage);

// 테스트 코드에서
MessageDto message = completableFuture.get(3, TimeUnit.SECONDS);

BlockingQueue는 여러 스레드에서 안전하게 접근할 수 있는 반면, CompletableFuture는 단일 결과를 비동기적으로 처리하는 데 유용합니다. 테스트의 복잡성과 요구사항에 따라 적절한 방식을 선택하면 됩니다.

WebSocket 클라이언트 구독 시 FrameHandler 설정

서버로부터 수신한 메시지를 처리하기 위해 StompFrameHandler를 구현한 핸들러를 설정합니다.

class DefaultStompFrameHandler implements StompFrameHandler {
    @Override
    public Type getPayloadType(StompHeaders headers) {
        return MessageDto.class;
    }

    @Override
    public void handleFrame(StompHeaders headers, Object payload) {
        blockingQueue.offer((MessageDto) payload);
    }
}

getPayloadType 메서드에서는 수신할 메시지의 타입을 지정했고, handleFrame 메서드에서는 수신한 메시지를 blockingQueue에 추가합니다.

Byte[]와 DTO 클래스 비교

초기에 테스트를 구현할 때, 컨버터를 적용하지 않아 Byte[]로 데이터를 처리했습니다. 현재는 컨버터로 변경하였지만 사용했던 경험이 있어 비교하여 적어봅니다.
StompFrameHandler의 getPayloadType 메서드에서 반환하는 타입을 byte[].class로 설정하면, 수신한 메시지를 바이트 배열로 처리하게 됩니다. 이 경우, 수동으로 바이트 배열을 원하는 객체로 변환해야 합니다.

@Override
public Type getPayloadType(StompHeaders headers) {
    return byte[].class;
}

@Override
public void handleFrame(StompHeaders headers, Object payload) {
    byte[] bytes = (byte[]) payload;
    // 바이트 배열을 문자열로 변환
    String message = new String(bytes, StandardCharsets.UTF_8);
    // JSON 문자열을 객체로 변환
    MessageDto messageDto = new ObjectMapper().readValue(message, MessageDto.class);
    blockingQueue.offer(messageDto);
}

반면, MessageDto.class로 설정하면 메시지 컨버터가 자동으로 JSON을 MessageDto 객체로 변환해주므로 코드가 더 간결해집니다. DTO 클래스를 직접 사용하는 것이 편리했습니다. 바이트 배열을 사용하는 것은 특별한 이유가 있을 때 사용하면 될 것 같습니다.

전체 코드 및 AssertJ를 사용한 검증

아래는 WebSocket 기능을 테스트하기 위한 전체 코드입니다.

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
public class WebSocketTest {

    @LocalServerPort
    private int port;

    private WebSocketStompClient stompClient;
    private BlockingQueue<MessageDto> blockingQueue;

    private String WEBSOCKET_URI;
    private final String WEBSOCKET_TOPIC = "/topic";

    @BeforeEach
    void setUp() {
        // 메시지 컨버터 설정
        MappingJackson2MessageConverter converter = new MappingJackson2MessageConverter();
        ObjectMapper objectMapper = converter.getObjectMapper().registerModule(new JavaTimeModule());

        // BlockingQueue 초기화
        blockingQueue = new LinkedBlockingQueue<>();

        // WebSocketStompClient 설정
        stompClient = new WebSocketStompClient(new SockJsClient(
                List.of(new WebSocketTransport(new StandardWebSocketClient()))));
        stompClient.setMessageConverter(converter);

        // WebSocket URI 설정
        WEBSOCKET_URI = "ws://localhost:" + port + "/ws";
    }

    @Test
    @DisplayName("WebSocket 메시지 전송 및 수신 테스트")
    void testWebSocketMessaging() throws Exception {
        Long roomId = 1L;
        String username = "user1";

        // WebSocket 연결 시 쿠키 설정
        HttpHeaders headers = new HttpHeaders();
        headers.add("Cookie", "login=" + username);

        // WebSocket 세션 연결
        StompSession session = stompClient
                .connectAsync(WEBSOCKET_URI + "?roomId=" + roomId, new WebSocketHttpHeaders(headers),
                        new StompSessionHandlerAdapter() {})
                        // 연결시도를 위한 최대 1초 대기
                .get(1, TimeUnit.SECONDS);

        // 특정 주제 구독
        session.subscribe(WEBSOCKET_TOPIC + "/room" + roomId, new DefaultStompFrameHandler());

        // 전송할 메시지 생성
        MessageDto messageToSend = new MessageDto("TEXT", "hi", "user1", "");

        // 메시지 전송
        session.send("/app/room" + roomId, messageToSend);

        // 메시지 수신 및 검증
        MessageDto receivedMessage = blockingQueue.poll(3, TimeUnit.SECONDS);

        assertThat(receivedMessage).isNotNull();
        assertThat(receivedMessage.getType()).as("메시지 타입 확인").isEqualTo(MessageType.TEXT);
        assertThat(receivedMessage.getContent()).as("메시지 내용 확인").isEqualTo(messageToSend.getContent());
        assertThat(receivedMessage.getSender()).as("메시지 발신자 확인").isEqualTo(messageToSend.getSender());
    }

    // StompFrameHandler 구현
    class DefaultStompFrameHandler implements StompFrameHandler {
        @Override
        public Type getPayloadType(StompHeaders headers) {
            return MessageDto.class;
        }

        @Override
        public void handleFrame(StompHeaders headers, Object payload) {
            blockingQueue.offer((MessageDto) payload);
        }
    }
}

결론

위의 테스트 코드를 통해 Spring Boot 애플리케이션에서 WebSocket 기능을 효과적으로 검증할 수 있었습니다. @SpringBootTest와 WebSocketStompClient, @ActiveProfile를 활용하여 통합 테스트를 구성하고, AssertJ를 사용하여 수신된 메시지의 내용을 검증함으로써 WebSocket 통신의 신뢰성을 확인할 수 있었습니다.