JUnit을 이용한 웹소켓 테스트

ksp7331·2023년 6월 30일

테스트 개요

이 글에서는 웹소켓 요청을 실제로 보낼 수 있게 해주는 WebSocketStompClient를 이용해 웹소켓을 테스트 한다. 따라서 테스트 환경에서 실제 서버를 동작시켜서 테스트를 진행한다.

테스트를 위한 설정

테스트 환경에서 실제 서버를 동작시키려면 다음과 같은 설정이 필요하다

// 8080이 서버의 포트번호가 된다.
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.DEFINED_PORT)
class SocketConnectionTest {
}

또는

// 포트번호가 랜덤이 되며, @LocalServerPort를 통해 포트번호를 불러올수 있다.
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class SocketConnectionTest {
	@LocalServerPort
    private int port;
}

웹소켓을 테스트할 때 WebSocketStompClient인스턴스를 사용하는데, 인스턴스를 생성할 때도 기본 설정이 필요하다. 이 기본설정은 웹소켓 설정에 따라 달라진다.

SockJS를 사용하지 않은 경우

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketBrokerConfig implements WebSocketMessageBrokerConfigurer {   

	@Override
	public void registerStompEndpoints(StompEndpointRegistry registry) {
		registry.addEndpoint("/ws/walk-logs")                
				.setAllowedOrigins("https://www.would-you-walk.com");
	}
}

WebSocketStompClient stompClient = new WebSocketStompClient(new StandardWebSocketClient());

SockJS를 사용한 경우

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketBrokerConfig implements WebSocketMessageBrokerConfigurer {   

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws/walk-logs")                
                .setAllowedOrigins("https://www.would-you-walk.com")
				.withSockJS();
    }
}

WebSocketStompClient stompClient = 
		new WebSocketStompClient(
				new SockJsClient(List.of(
						new WebSocketTransport(
								new StandardWebSocketClient()
						)
				)
		)
);

WebSocketStompClient 인스턴스를 생성 후, 직렬화, 역직렬화를 위한 converter를 지정해야 한다

stompClient.setMessageConverter(new MappingJackson2MessageConverter());

이 때, 역직렬화시 LocalDateTime타입으로의 변환이 필요하다면, JavaTimeModule이 필요하다. 또한 역직렬화시 모든 파라미터를 매개변수로 받는 생성자(Lombok의 @AllArgsConstructor)를 이용하려면 ParameterNamesModule을 추가해야 한다.

MappingJackson2MessageConverter messageConverter = new MappingJackson2MessageConverter();
ObjectMapper objectMapper = messageConverter.getObjectMapper();
objectMapper.registerModules(new JavaTimeModule(), new ParameterNamesModule());
stompClient.setMessageConverter(messageConverter);

웹소켓 연결하기

웹소켓 연결은 WebSocketStompClient의 connect()를 사용해서 한다. connect()는 여러 형태가 메서드 오버로딩되어 있으므로 필요한것을 사용하면 된다.

// (1)
ListenableFuture<StompSession> connect(String url, StompSessionHandler handler, Object... uriVars)
// (2)
ListenableFuture<StompSession> connect(String url, @Nullable WebSocketHttpHeaders handshakeHeaders,
			StompSessionHandler handler, Object... uriVariables)
// (3)
ListenableFuture<StompSession> connect(String url, @Nullable WebSocketHttpHeaders handshakeHeaders,
			@Nullable StompHeaders connectHeaders, StompSessionHandler handler, Object... uriVariables)
// (4)
ListenableFuture<StompSession> connect(URI url, @Nullable WebSocketHttpHeaders handshakeHeaders,
			@Nullable StompHeaders connectHeaders, StompSessionHandler sessionHandler)

url

url는 ws 프로토콜을 사용해야 한다.

this.url = String.format("ws://localhost:%d/ws/walk-logs", port);

handler

StompSessionHandler 타입의 인스턴스가 필요한데, 추상클래스인 StompSessionHandlerAdapter의 익명클래스를 사용한다.
```
stompClient.connect(url, new StompSessionHandlerAdapter(){});
```
handshakeHeaders

웹소켓 핸드셰이크 요청에 헤더를 추가하는 매개변수이다. 이를 통해 웹소켓 연결시 원하는 헤더를 포함시킬 수 있지만.. 사용하지 않는것이 좋다. 이유는 일반적으로 클라이언트에서 웹소켓 핸드셰이크 요청을 할때 헤더를 추가하는것이 허용되지 않기 때문이다. 자세한 내용은 아래 블로그에 나와있다.

Socket 인증 with API Gateway + Refresh JWT
connectHeaders

CONNECT 프레임의 메시지에 헤더를 추가한다. handshakeHeaders 없이 connectHeaders만 매개변수에 포함시킬 수 없기 때문에 connectHeaders를 사용하려면 handshakeHeaders자리에 빈 인스턴스를 넣어야 한다.
```
stompClient.connect(url, new WebSocketHttpHeaders(), stompHeaders, new StompSessionHandlerAdapter(){});
```

return type

connect()는 비동기적으로 동작하기 때문에 결과값이 ListenableFuture<StompSession>로 나온다. 여기서 추가작업을 위해서는 StompSession이 필요하기 때문에, get()을 사용해서 StompSession을 추출해야 한다.

// stompClient.connect()가 완료되고 stompSession를 반환할 때까지 기다린다. 단, 2초가 지나면 TimeOutException이 발생한다.
this.stompSession = stompClient.connect(url, new WebSocketHttpHeaders(), stompHeaders, new StompSessionHandlerAdapter() {
        }).get(2, TimeUnit.SECONDS);

subscribe()를 사용한다. 매개변수에서connect()와의 차이점이 있다면 두번째 매개변수로 StompFrameHandler타입의 인스턴스를 사용한다. StompFrameHandler는 StompSessionHandler의 상위 인터페이스이다.

stompSession.subscribe("/topic/1", new StompSessionHandlerAdapter(){});

메시지 전송하기

send()를 사용한다.

stompSession.send("/topic/1", pub);

웹소켓이 동작하는 스레드는 테스트가 동작하는 스레드와 다르다.
아래쪽이 웹소켓이 동작하는 스레드이다.

테스트가 동작하는 스레드는 더이상 실행할 코드가 없으면, 웹소켓 스레드를 기다리지 않고 그대로 테스트를 종료시킨다. 따라서 메시지를 보낸 후 보낸 메시지를 받기 위해서는 웹소켓 스레드의 동작이 끝날때까지 기다려주는 코드가 필요하다.

이를 위해 CompletableFuture를 사용한다.

// 타입 매개변수에는 JSON 메시지를 역직렬화 했을때의 타입을 적으면 된다.
CompletableFuture<CoordinateControllerDTO.Sub> subscribeFuture = new CompletableFuture<>();

// 중략

stompSession.subscribe("topic", new StompFrameHandler() {
						// 여기서 지정한 타입으로 payload를 역직렬화한다.
            @Override
            public Type getPayloadType(StompHeaders headers) {                
                return CoordinateControllerDTO.Sub.class;
            }

						// 여기서 메시지를 가져올 수 있다.
            @Override
            public void handleFrame(StompHeaders headers, Object payload) {                
                subscribeFuture.complete((CoordinateControllerDTO.Sub) payload);
            }
        });

// subscribeFuture인스턴스에 값이 들어올때 까지 최대 3초 기다린다.
CoordinateControllerDTO.Sub sub = subscribeFuture.get(3, TimeUnit.SECONDS);