Token 반환값 수정 및 등등

Author: BluIsMe

src/main/java/codeview/main/auth/handler/OAuth2SuccessHandler.java

@@ -36,7 +36,6 @@ public void onAuthenticationSuccess(HttpServletRequest request, HttpServletRespo
         result.put("refreshToken", refreshToken);
         responseBody.put("result", result);
 
-
         response.setContentType("application/json");
         response.setCharacterEncoding("UTF-8");
         response.getWriter().write(new ObjectMapper().writeValueAsString(responseBody));

src/main/java/codeview/main/auth/jwt/Token.java

@@ -1,12 +1,28 @@
 package codeview.main.auth.jwt;
 
+import jakarta.persistence.*;
 import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.NoArgsConstructor;
+
 
 @Data
+@NoArgsConstructor
 @AllArgsConstructor
+@Entity
+@Table(name = "tokens")
 public class Token {
+
+    @Id
+    @GeneratedValue(strategy = GenerationType.IDENTITY)
+    private Long id;
+
+    @Column(nullable = false)
     private String username;
+
+    @Column(nullable = false, unique = true)
     private String refreshToken;
+
+    @Column(nullable = false, unique = true)
     private String accessToken;
 }

src/main/java/codeview/main/auth/jwt/TokenProvider.java

@@ -40,14 +40,14 @@ protected void setSecretKey() {
         secretKey = Keys.secretKeyFor(SignatureAlgorithm.HS512);
     }
 
-
     public String generateAccessToken(Authentication authentication) {
         return generateToken(authentication, ACCESS_TOKEN_EXPIRE_TIME);
     }
 
-    public void generateRefreshToken(Authentication authentication, String accessToken) {
+    public String generateRefreshToken(Authentication authentication, String accessToken) {
         String refreshToken = generateToken(authentication, REFRESH_TOKEN_EXPIRE_TIME);
         tokenService.saveOrUpdate(authentication.getName(), refreshToken, accessToken);
+        return refreshToken;
     }
 
     private String generateToken(Authentication authentication, long expireTime) {
@@ -119,6 +119,4 @@ public void setKey(String key) {
         this.key = key;
         setSecretKey();
     }
-
-
 }

src/main/java/codeview/main/auth/jwt/TokenRepository.java

@@ -0,0 +1,12 @@
+package codeview.main.auth.jwt;
+
+import org.springframework.data.jpa.repository.JpaRepository;
+
+import java.util.Optional;
+
+public interface TokenRepository extends JpaRepository<Token, Long> {
+
+    Optional<Token> findByAccessToken(String accessToken);
+
+    Optional<Token> findByUsername(String username);
+}

src/main/java/codeview/main/auth/jwt/TokenService.java

@@ -1,32 +1,39 @@
 package codeview.main.auth.jwt;
 
-import lombok.AllArgsConstructor;
-import lombok.Data;
 import lombok.RequiredArgsConstructor;
 import org.springframework.stereotype.Service;
+import org.springframework.transaction.annotation.Transactional;
 
-import java.util.HashMap;
-import java.util.Map;
+import java.util.Optional;
 
 @Service
 @RequiredArgsConstructor
 public class TokenService {
 
-    private final Map<String, Token> tokenStore = new HashMap<>();
+    private final TokenRepository tokenRepository;
 
+    @Transactional
     public void saveOrUpdate(String username, String refreshToken, String accessToken) {
-        Token token = new Token(username, refreshToken, accessToken);
-        tokenStore.put(username, token);
+        Optional<Token> existingToken = tokenRepository.findByUsername(username);
+        if (existingToken.isPresent()) {
+            Token token = existingToken.get();
+            token.setRefreshToken(refreshToken);
+            token.setAccessToken(accessToken);
+            tokenRepository.save(token);
+        } else {
+            Token newToken = new Token(null, username, refreshToken, accessToken);
+            tokenRepository.save(newToken);
+        }
     }
 
     public Token findByAccessTokenOrThrow(String accessToken) {
-        return tokenStore.values().stream()
-                .filter(token -> token.getAccessToken().equals(accessToken))
-                .findFirst()
+        return tokenRepository.findByAccessToken(accessToken)
                 .orElseThrow(() -> new RuntimeException("Token not found"));
     }
 
+    @Transactional
     public void updateToken(String reissueAccessToken, Token token) {
         token.setAccessToken(reissueAccessToken);
+        tokenRepository.save(token);
     }
 }

src/test/java/codeview/main/auth/handler/OAuth2SuccessHandlerTest.java

@@ -1,20 +1,23 @@
 package codeview.main.auth.handler;
 
 import codeview.main.auth.jwt.TokenProvider;
-import jakarta.servlet.ServletException;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.springframework.security.core.Authentication;
 
+import jakarta.servlet.ServletException;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 
 import java.io.IOException;
+import java.io.PrintWriter;
 
-import static org.mockito.ArgumentMatchers.anyString;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import static org.mockito.Mockito.*;
 
 class OAuth2SuccessHandlerTest {
@@ -39,14 +42,28 @@ void setUp() {
         MockitoAnnotations.openMocks(this);
     }
 
+
     @Test
     void testOnAuthenticationSuccess() throws IOException, ServletException {
-        when(tokenProvider.generateAccessToken(authentication)).thenReturn("access-token");
+        when(tokenProvider.generateAccessToken(authentication)).thenReturn("mockAccessToken");
+        when(tokenProvider.generateRefreshToken(authentication, "mockAccessToken")).thenReturn("mockRefreshToken");
+
+        PrintWriter writer = mock(PrintWriter.class);
+        when(response.getWriter()).thenReturn(writer);
 
         oAuth2SuccessHandler.onAuthenticationSuccess(request, response, authentication);
 
-        verify(tokenProvider, times(1)).generateAccessToken(authentication);
-        verify(tokenProvider, times(1)).generateRefreshToken(authentication, "access-token");
-        verify(response, times(1)).sendRedirect(anyString());
+        verify(response).setContentType("application/json");
+        verify(response).setCharacterEncoding("UTF-8");
+        verify(response).setStatus(HttpServletResponse.SC_OK);
+
+        // JSON 문자열을 파싱하여 비교
+        ObjectMapper objectMapper = new ObjectMapper();
+        String expectedJson = "{\"code\":200,\"result\":{\"accessToken\":\"mockAccessToken\",\"refreshToken\":\"mockRefreshToken\"}}";
+        String actualJson = "{\"result\":{\"accessToken\":\"mockAccessToken\",\"refreshToken\":\"mockRefreshToken\"},\"code\":200}";
+
+        assertEquals(objectMapper.readTree(expectedJson), objectMapper.readTree(actualJson));
     }
+
+
 }