qingsongping-AIDevelop/mydreamjob-server/src/test/java/com/mydreamjob/messaging/realtime/ChatSocketAuthTest.java

package com.mydreamjob.messaging.realtime;

import com.mydreamjob.auth.*;
import org.junit.jupiter.api.*;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.messaging.simp.stomp.*;
import java.time.*;
import java.util.*;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;

class ChatSocketAuthTest {
    private final Instant now = Instant.parse("2026-10-14T08:00:00Z");
    private JwtTokenProvider tokens;
    private AuthSessionService sessions;
    private UserRepository users;
    private ChatSocketAuth guard;
    private AccessTokenClaims claims;

    @BeforeEach void setup() {
        tokens = mock(JwtTokenProvider.class); sessions = mock(AuthSessionService.class);
        users = mock(UserRepository.class);
        guard = new ChatSocketAuth(tokens, sessions, users,
                mock(com.mydreamjob.company.CompanyRepository.class), Clock.fixed(now, ZoneOffset.UTC));
        claims = new AccessTokenClaims(1L, "session", "token", "STUDENT", now.minusSeconds(1), now.plusSeconds(60));
        when(tokens.parseAndValidateAccessToken("valid", now)).thenReturn(claims);
        when(sessions.isSessionActive(1L, "session", now)).thenReturn(true);
        when(users.findByIdWithRoles(1L)).thenReturn(Optional.of(User.builder().id(1L).build()));
    }
    private StompHeaderAccessor frame(StompCommand command, Map<String, Object> attrs) {
        var frame = StompHeaderAccessor.create(command);
        frame.setSessionId("ws-1"); frame.setSessionAttributes(attrs); frame.setLeaveMutable(true);
        return frame;
    }
    private void send(StompHeaderAccessor frame) {
        guard.preSend(MessageBuilder.createMessage(new byte[0], frame.getMessageHeaders()), null);
    }
    private Map<String, Object> connect() {
        var attrs = new HashMap<String, Object>(); var frame = frame(StompCommand.CONNECT, attrs);
        frame.setNativeHeader("Authorization", "Bearer valid"); send(frame); return attrs;
    }
    @Test void connectAuthenticatesAndRemovesRawToken() {
        var frame = frame(StompCommand.CONNECT, new HashMap<>());
        frame.setNativeHeader("Authorization", "Bearer valid"); send(frame);
        assertEquals("1", frame.getUser().getName());
        assertNull(frame.getFirstNativeHeader("Authorization"));
    }
    @Test void anonymousConnectRejected() {
        assertThrows(Exception.class, () -> send(frame(StompCommand.CONNECT, new HashMap<>())));
    }
    @Test void revokedSessionRejected() {
        when(sessions.isSessionActive(1L, "session", now)).thenReturn(false);
        assertThrows(Exception.class, this::connect);
    }
    @Test void expiredTokenRejected() {
        when(tokens.parseAndValidateAccessToken("valid", now)).thenReturn(
                new AccessTokenClaims(1L, "session", "token", "STUDENT", now.minusSeconds(60), now));
        assertThrows(Exception.class, this::connect);
    }
    @Test void unassignedIdentityRejected() {
        when(users.findByIdWithRoles(1L)).thenReturn(Optional.of(User.builder().role(User.Role.UNASSIGNED).build()));
        assertThrows(Exception.class, this::connect);
    }
    @Test void ownFixedQueueAllowed() {
        var frame = frame(StompCommand.SUBSCRIBE, connect());
        frame.setDestination("/user/queue/chat-events"); assertDoesNotThrow(() -> send(frame));
    }
    @Test void otherUserAndBrokerQueuesRejected() {
        var attrs = connect();
        for (String path : List.of("/user/2/queue/chat-events", "/queue/chat-events", "/topic/all")) {
            var frame = frame(StompCommand.SUBSCRIBE, attrs); frame.setDestination(path);
            assertThrows(Exception.class, () -> send(frame));
        }
    }
    @Test void socketSendRejected() {
        var frame = frame(StompCommand.SEND, connect()); frame.setDestination("/app/chat");
        assertThrows(Exception.class, () -> send(frame));
    }
    @Test void databaseFailureFailsClosed() {
        when(users.findByIdWithRoles(1L)).thenThrow(new IllegalStateException("private detail"));
        var error = assertThrows(Exception.class, this::connect);
        assertEquals("Chat connection rejected", error.getMessage());
    }
    @Test void repeatedConnectRejected() {
        var frame = frame(StompCommand.CONNECT, connect()); frame.setNativeHeader("Authorization", "Bearer valid");
        assertThrows(Exception.class, () -> send(frame));
    }
    @Test void unverifiedCompanyRejected() {
        when(users.findByIdWithRoles(1L)).thenReturn(Optional.of(User.builder().id(1L)
                .role(User.Role.COMPANY).roles(Set.of(User.Role.COMPANY)).build()));
        assertThrows(Exception.class, this::connect);
    }
    @Test void revokedSessionCannotSubscribeAgain() {
        var attrs = connect(); when(sessions.isSessionActive(1L, "session", now)).thenReturn(false);
        var frame = frame(StompCommand.SUBSCRIBE, attrs); frame.setDestination("/user/queue/chat-events");
        assertThrows(Exception.class, () -> send(frame));
    }
    @Test void subscriptionFloodRejected() {
        var attrs = connect(); var frame = frame(StompCommand.SUBSCRIBE, attrs);
        frame.setDestination("/user/queue/chat-events");
        for (int i = 0; i < 119; i++) send(frame);
        assertThrows(Exception.class, () -> send(frame));
    }
    @Test void outboundChecksRevocationBeforeDelivery() {
        connect();
        var headers = org.springframework.messaging.simp.SimpMessageHeaderAccessor.create(
                org.springframework.messaging.simp.SimpMessageType.MESSAGE);
        headers.setSessionId("ws-1");
        var message = MessageBuilder.createMessage(new byte[0], headers.getMessageHeaders());
        assertNotNull(guard.outbound().preSend(message, null));
        when(sessions.isSessionActive(1L, "session", now)).thenReturn(false);
        assertNull(guard.outbound().preSend(message, null));
    }
}