tuanwei-java/src/test/java/com/gxwebsoft/ai/config/ChatStreamRegistryTest.java

package com.gxwebsoft.ai.config;

import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.function.Consumer;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * HTTP 流式通道（SSE）是生产上唯一能用的推送路径：网关在 {@code /api} 前丢掉
 * WebSocket 的 {@code Upgrade} 头，WS 握手永远到不了 Tomcat。
 *
 * <p>因此这张登记表的几个行为是硬约束：没人接时静默、有人接时逐字转发、
 * 重复登记后者覆盖前者、写失败要把死掉的 sink 摘掉
 * （否则后续消息会一直往已经断开的长连接里写）。
 *
 * <p>每个用例用独立 userId，避免静态表互相串味，也就不需要清理。
 */
class ChatStreamRegistryTest {

    private static String newUser() {
        return "test-" + UUID.randomUUID();
    }

    @Test
    void 没有登记通道时返回false() {
        assertFalse(ChatStreamRegistry.emit(newUser(), "{}"));
    }

    @Test
    void 登记后能收到原文() {
        String user = newUser();
        List<String> received = new ArrayList<>();
        ChatStreamRegistry.bind(user, received::add);

        assertTrue(ChatStreamRegistry.emit(user, "{\"answer\":\"你好\"}"));
        assertEquals(Collections.singletonList("{\"answer\":\"你好\"}"), received);
    }

    @Test
    void 重复登记时后者覆盖前者() {
        String user = newUser();
        List<String> first = new ArrayList<>();
        List<String> second = new ArrayList<>();
        ChatStreamRegistry.bind(user, first::add);
        ChatStreamRegistry.bind(user, second::add);

        ChatStreamRegistry.emit(user, "x");

        assertTrue(first.isEmpty(), "被覆盖的旧连接不该再收到消息");
        assertEquals(Collections.singletonList("x"), second);
    }

    @Test
    void 解绑只摘自己的那个通道() {
        String user = newUser();
        List<String> first = new ArrayList<>();
        List<String> second = new ArrayList<>();
        Consumer<String> firstSink = first::add;
        ChatStreamRegistry.bind(user, firstSink);
        ChatStreamRegistry.bind(user, second::add);

        // 旧连接晚到的 finally 块：不能把新连接顶掉
        ChatStreamRegistry.unbind(user, firstSink);

        assertTrue(ChatStreamRegistry.emit(user, "still-alive"));
        assertEquals(Collections.singletonList("still-alive"), second);
    }

    @Test
    void 写入失败后摘掉通道且不向上抛() {
        String user = newUser();
        ChatStreamRegistry.bind(user, payload -> {
            throw new IllegalStateException("客户端已断开");
        });

        assertFalse(ChatStreamRegistry.emit(user, "boom"));
        // 失败过的 sink 已被摘掉，下一次 emit 直接返回 false
        assertFalse(ChatStreamRegistry.emit(user, "boom"));
    }

    @Test
    void 没有WebSocket会话时推送落到流式通道() throws IOException {
        String user = newUser();
        List<String> received = new ArrayList<>();
        ChatStreamRegistry.bind(user, received::add);

        // 生产就是这样：SESSIONS 里没有这个用户，send 必须改投 SSE
        WebSocketServer.send(user, "{\"answer\":\"fallback\"}");

        assertEquals(Collections.singletonList("{\"answer\":\"fallback\"}"), received);
    }
}