Skip to content

Commit 71234e9

Browse files
jueunk617jueunk617
committed
Refactor: WebSocketMessageController에서 인증된 Principal 정보를 사용하도록 수정
1 parent 686bf28 commit 71234e9

File tree

3 files changed

+59
-51
lines changed

3 files changed

+59
-51
lines changed

‎src/main/java/com/back/global/websocket/controller/WebSocketMessageController.java‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
package com.back.global.websocket.controller;
22

33
import com.back.global.exception.CustomException;
4-
import com.back.global.websocket.dto.HeartbeatMessage;
4+
import java.security.Principal;
5+
import com.back.global.security.user.CustomUserDetails;
56
import com.back.global.websocket.service.WebSocketSessionManager;
67
import com.back.global.websocket.util.WebSocketErrorHelper;
78
import lombok.RequiredArgsConstructor;
89
import lombok.extern.slf4j.Slf4j;
910
import org.springframework.messaging.handler.annotation.MessageMapping;
1011
import org.springframework.messaging.handler.annotation.Payload;
1112
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
13+
import org.springframework.security.core.Authentication;
1214
import org.springframework.stereotype.Controller;
1315

1416
@Slf4j
@@ -21,16 +23,18 @@ public class WebSocketMessageController {
2123

2224
// Heartbeat 처리
2325
@MessageMapping("/heartbeat")
24-
public void handleHeartbeat(@Payload HeartbeatMessage message,
26+
public void handleHeartbeat(Principal principal,
2527
SimpMessageHeaderAccessor headerAccessor) {
2628
try {
27-
if (message.userId() != null) {
28-
// TTL 10분으로 연장
29-
sessionManager.updateLastActivity(message.userId());
30-
log.debug("Heartbeat 처리 완료 - 사용자: {}", message.userId());
29+
// Principal에서 인증된 사용자 정보 추출
30+
if (principal instanceof Authentication auth && auth.getPrincipal() instanceof CustomUserDetails userDetails) {
31+
Long userId = userDetails.getUserId();
32+
33+
sessionManager.updateLastActivity(userId);
34+
log.debug("Heartbeat 처리 완료 - 사용자: {}", userId);
3135
} else {
32-
log.warn("유효하지 않은 Heartbeat 메시지 수신: userId가 null");
33-
errorHelper.sendInvalidRequestError(headerAccessor.getSessionId(), "사용자 ID가 필요합니다");
36+
log.warn("인증되지 않은 Heartbeat 요청: {}", headerAccessor.getSessionId());
37+
errorHelper.sendUnauthorizedError(headerAccessor.getSessionId());
3438
}
3539
} catch (CustomException e) {
3640
log.error("Heartbeat 처리 실패: {}", e.getMessage());
@@ -43,12 +47,14 @@ public void handleHeartbeat(@Payload HeartbeatMessage message,
4347

4448
// 사용자 활동 신호 처리
4549
@MessageMapping("/activity")
46-
public void handleActivity(@Payload HeartbeatMessage message,
50+
public void handleActivity(Principal principal,
4751
SimpMessageHeaderAccessor headerAccessor) {
4852
try {
49-
if (message.userId() != null) {
50-
sessionManager.updateLastActivity(message.userId());
51-
log.debug("사용자 활동 신호 처리 완료 - 사용자: {}", message.userId());
53+
if (principal instanceof Authentication auth && auth.getPrincipal() instanceof CustomUserDetails userDetails) {
54+
Long userId = userDetails.getUserId();
55+
56+
sessionManager.updateLastActivity(userId);
57+
log.debug("사용자 활동 신호 처리 완료 - 사용자: {}", userId);
5258
} else {
5359
log.warn("유효하지 않은 활동 신호: userId가 null");
5460
errorHelper.sendInvalidRequestError(headerAccessor.getSessionId(), "사용자 ID가 필요합니다");

‎src/main/java/com/back/global/websocket/dto/HeartbeatMessage.java‎

Lines changed: 0 additions & 5 deletions
This file was deleted.

‎src/test/java/com/back/global/websocket/controller/WebSocketMessageControllerTest.java‎

Lines changed: 41 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import com.back.global.exception.CustomException;
44
import com.back.global.exception.ErrorCode;
5-
import com.back.global.websocket.dto.HeartbeatMessage;
5+
import com.back.global.security.user.CustomUserDetails;
66
import com.back.global.websocket.service.WebSocketSessionManager;
77
import com.back.global.websocket.util.WebSocketErrorHelper;
88
import org.junit.jupiter.api.BeforeEach;
@@ -14,6 +14,9 @@
1414
import org.mockito.Mock;
1515
import org.mockito.junit.jupiter.MockitoExtension;
1616
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
17+
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
18+
19+
import java.security.Principal;
1720

1821
import static org.mockito.ArgumentMatchers.*;
1922
import static org.mockito.Mockito.*;
@@ -33,64 +36,68 @@ class WebSocketMessageControllerTest {
3336

3437
private SimpMessageHeaderAccessor headerAccessor;
3538
private Long userId;
36-
private Long roomId;
3739
private String sessionId;
3840

3941
@BeforeEach
4042
void setUp() {
4143
userId = 10L;
42-
roomId = 1L;
4344
sessionId = "test-session-id";
4445

4546
headerAccessor = mock(SimpMessageHeaderAccessor.class);
4647
lenient().when(headerAccessor.getSessionId()).thenReturn(sessionId);
4748
}
4849

50+
// 테스트용 Mock Principal 객체를 생성하는 헬퍼 메소드
51+
private Principal createMockPrincipal(Long userId) {
52+
CustomUserDetails mockUserDetails = mock(CustomUserDetails.class);
53+
when(mockUserDetails.getUserId()).thenReturn(userId);
54+
55+
return new UsernamePasswordAuthenticationToken(mockUserDetails, null, null);
56+
}
57+
4958
@Nested
5059
@DisplayName("Heartbeat 처리")
5160
class HandleHeartbeatTest {
5261

5362
@Test
54-
@DisplayName("정상 - Heartbeat 처리")
63+
@DisplayName("성공 - 인증된 사용자의 Heartbeat 처리")
5564
void t1() {
5665
// given
57-
HeartbeatMessage message = new HeartbeatMessage(userId);
66+
Principal mockPrincipal = createMockPrincipal(userId);
5867
doNothing().when(sessionManager).updateLastActivity(userId);
5968

6069
// when
61-
controller.handleHeartbeat(message, headerAccessor);
70+
controller.handleHeartbeat(mockPrincipal, headerAccessor);
6271

6372
// then
6473
verify(sessionManager).updateLastActivity(userId);
65-
verify(errorHelper, never()).sendInvalidRequestError(anyString(), anyString());
66-
verify(errorHelper, never()).sendCustomExceptionToUser(anyString(), any());
67-
verify(errorHelper, never()).sendGenericErrorToUser(anyString(), any(), anyString());
74+
verifyNoInteractions(errorHelper);
6875
}
6976

7077
@Test
71-
@DisplayName("실패 - userId가 null")
78+
@DisplayName("실패 - 인증 정보가 없는 경우")
7279
void t2() {
7380
// given
74-
HeartbeatMessage message = new HeartbeatMessage(null);
81+
Principal principal = null;
7582

7683
// when
77-
controller.handleHeartbeat(message, headerAccessor);
84+
controller.handleHeartbeat(principal, headerAccessor);
7885

7986
// then
8087
verify(sessionManager, never()).updateLastActivity(any());
81-
verify(errorHelper).sendInvalidRequestError(sessionId, "사용자 ID가 필요합니다");
88+
verify(errorHelper).sendUnauthorizedError(sessionId);
8289
}
8390

8491
@Test
8592
@DisplayName("실패 - CustomException 발생")
8693
void t3() {
8794
// given
88-
HeartbeatMessage message = new HeartbeatMessage(userId);
95+
Principal mockPrincipal = createMockPrincipal(userId);
8996
CustomException exception = new CustomException(ErrorCode.BAD_REQUEST);
9097
doThrow(exception).when(sessionManager).updateLastActivity(userId);
9198

9299
// when
93-
controller.handleHeartbeat(message, headerAccessor);
100+
controller.handleHeartbeat(mockPrincipal, headerAccessor);
94101

95102
// then
96103
verify(sessionManager).updateLastActivity(userId);
@@ -101,12 +108,12 @@ void t3() {
101108
@DisplayName("실패 - 일반 Exception 발생")
102109
void t4() {
103110
// given
104-
HeartbeatMessage message = new HeartbeatMessage(userId);
111+
Principal mockPrincipal = createMockPrincipal(userId);
105112
RuntimeException exception = new RuntimeException("예상치 못한 오류");
106113
doThrow(exception).when(sessionManager).updateLastActivity(userId);
107114

108115
// when
109-
controller.handleHeartbeat(message, headerAccessor);
116+
controller.handleHeartbeat(mockPrincipal, headerAccessor);
110117

111118
// then
112119
verify(sessionManager).updateLastActivity(userId);
@@ -123,46 +130,46 @@ void t4() {
123130
class HandleActivityTest {
124131

125132
@Test
126-
@DisplayName("정상 - 활동 신호 처리")
127-
void t13() {
133+
@DisplayName("성공 - 인증된 사용자의 활동 신호 처리")
134+
void t1() {
128135
// given
129-
HeartbeatMessage message = new HeartbeatMessage(userId);
136+
Principal mockPrincipal = createMockPrincipal(userId);
130137
doNothing().when(sessionManager).updateLastActivity(userId);
131138

132139
// when
133-
controller.handleActivity(message, headerAccessor);
140+
controller.handleActivity(mockPrincipal, headerAccessor);
134141

135142
// then
136143
verify(sessionManager).updateLastActivity(userId);
137-
verify(errorHelper, never()).sendInvalidRequestError(anyString(), anyString());
138-
verify(errorHelper, never()).sendCustomExceptionToUser(anyString(), any());
139-
verify(errorHelper, never()).sendGenericErrorToUser(anyString(), any(), anyString());
144+
verifyNoInteractions(errorHelper);
140145
}
141146

142147
@Test
143-
@DisplayName("실패 - userId가 null")
144-
void t14() {
148+
@DisplayName("실패 - 인증 정보가 없는 경우")
149+
void t2() {
145150
// given
146-
HeartbeatMessage message = new HeartbeatMessage(null);
151+
Principal principal = null;
147152

148153
// when
149-
controller.handleActivity(message, headerAccessor);
154+
controller.handleActivity(principal, headerAccessor);
150155

151156
// then
152157
verify(sessionManager, never()).updateLastActivity(any());
158+
159+
// handleActivity의 else 블록에 맞춰 검증 로직 수정
153160
verify(errorHelper).sendInvalidRequestError(sessionId, "사용자 ID가 필요합니다");
154161
}
155162

156163
@Test
157164
@DisplayName("실패 - CustomException 발생")
158-
void t15() {
165+
void t3() {
159166
// given
160-
HeartbeatMessage message = new HeartbeatMessage(userId);
167+
Principal mockPrincipal = createMockPrincipal(userId);
161168
CustomException exception = new CustomException(ErrorCode.BAD_REQUEST);
162169
doThrow(exception).when(sessionManager).updateLastActivity(userId);
163170

164171
// when
165-
controller.handleActivity(message, headerAccessor);
172+
controller.handleActivity(mockPrincipal, headerAccessor);
166173

167174
// then
168175
verify(sessionManager).updateLastActivity(userId);
@@ -171,14 +178,14 @@ void t15() {
171178

172179
@Test
173180
@DisplayName("실패 - 일반 Exception 발생")
174-
void t16() {
181+
void t4() {
175182
// given
176-
HeartbeatMessage message = new HeartbeatMessage(userId);
183+
Principal mockPrincipal = createMockPrincipal(userId);
177184
RuntimeException exception = new RuntimeException("예상치 못한 오류");
178185
doThrow(exception).when(sessionManager).updateLastActivity(userId);
179186

180187
// when
181-
controller.handleActivity(message, headerAccessor);
188+
controller.handleActivity(mockPrincipal, headerAccessor);
182189

183190
// then
184191
verify(sessionManager).updateLastActivity(userId);

0 commit comments

Comments
 (0)