如何在Go中测试Fiber框架的SSE端点客户端断开行为？

我正在为基于Fiber框架实现的Server-Sent Events（SSE）端点编写Go测试，目标是验证即便服务器被设计为无限保持连接，客户端仍能成功断开。目前的基础测试代码会陷入阻塞，因为app.Test(req, -1)会无限等待响应完成，而SSE长连接不会主动终止。现有测试代码如下：

func Test_WatcherStream_Returns200(t *testing.T) {
    entityID := uuid.New().String()

    req := httptest.NewRequest("GET", fmt.Sprintf("/watchers/%s", entityID), nil)
    resp, err := app.Test(req, -1)
    
    assert.NoError(t, err)
    assert.Equal(t, http.StatusOK, resp.StatusCode)

    err = resp.Body.Close()
    assert.NoError(t, err)
}

解决方案 #

• 核心问题：app.Test的第二个参数设为-1会让测试无限等待响应结束，但SSE是长连接，服务器不会主动关闭。需要通过上下文取消或超时控制模拟客户端断开，并让服务器正确处理连接终止。
• 改造步骤：
  1. 创建带可取消上下文的请求，将上下文注入到Fiber的请求中（Fiber会尊重请求的上下文取消信号）。
  2. 在获取响应并验证状态码后，主动关闭响应体（模拟客户端断开），同时取消上下文，触发服务器端的连接终止逻辑。
  3. 使用goroutine配合sync.WaitGroup或超时机制，确保测试能正常结束，不会无限阻塞。

改造后的测试代码

func Test_WatcherStream_ClientDisconnect(t *testing.T) {
    entityID := uuid.New().String()

    // 创建可取消上下文
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    // 注入上下文到请求
    req := httptest.NewRequest("GET", fmt.Sprintf("/watchers/%s", entityID), nil)
    req = req.WithContext(ctx)

    var resp *http.Response
    var testErr error

    // 启动goroutine执行请求，避免主线程阻塞
    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        defer wg.Done()
        // 设置超时防止极端情况卡住
        resp, testErr = app.Test(req, 5*time.Second)
    }()

    // 等待连接建立并返回响应头
    time.Sleep(100 * time.Millisecond)

    // 验证响应状态
    assert.NoError(t, testErr)
    assert.Equal(t, http.StatusOK, resp.StatusCode)

    // 模拟客户端断开：关闭响应体
    assert.NoError(t, resp.Body.Close())

    // 取消上下文，通知服务器终止连接
    cancel()

    // 等待goroutine完成
    wg.Wait()
}

服务器端SSE处理配合逻辑

要让测试生效，你的SSE端点处理函数需要监听请求上下文的取消信号，当客户端断开时及时退出：

func watcherStream(c *fiber.Ctx) error {
    c.Set("Content-Type", "text/event-stream")
    c.Set("Cache-Control", "no-cache")
    c.Set("Connection", "keep-alive")

    ctx := c.Context()
    for {
        select {
        case <-ctx.Done():
            // 客户端断开，退出循环释放资源
            return nil
        default:
            // 发送SSE心跳或事件
            _, err := c.WriteString("data: ping\n\n")
            if err != nil {
                return err
            }
            // 刷新响应缓冲区
            c.Context().Response.BodyFlush()
            time.Sleep(1 * time.Second)
        }
    }
}

内容的提问来源于stack exchange，提问作者Jay