//go:build e2e // websocket_e2e_test.go - WebSocket E2E 测试 // // 测试 lolly WebSocket 代理功能:连接建立、消息传递、并发连接等。 // // 作者:xfy package e2e import ( "context" "fmt" "net/http" "net/http/httptest" "testing" "time" "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "rua.plus/lolly/internal/e2e/testutil" ) // wsEchoHandler WebSocket Echo 处理器。 func wsEchoHandler(w http.ResponseWriter, r *http.Request) { upgrader := websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { return } defer conn.Close() for { messageType, data, err := conn.ReadMessage() if err != nil { return } if err := conn.WriteMessage(messageType, data); err != nil { return } } } // TestE2EWebSocketBasic 测试基础 WebSocket 代理。 // // 验证 WebSocket 连接可以成功建立和消息传递。 func TestE2EWebSocketBasic(t *testing.T) { t.Parallel() // 创建本地 WebSocket Echo 服务器 server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() // 将 http:// 替换为 ws:// wsURL := "ws" + server.URL[4:] t.Logf("WebSocket Echo server: %s", wsURL) // 创建 WebSocket 客户端 client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送消息 testMessage := "Hello, WebSocket!" err = client.Send(testMessage) require.NoError(t, err, "Failed to send message") // 接收响应 response, err := client.Receive() require.NoError(t, err, "Failed to receive message") // Echo 服务器应该返回相同的消息 assert.Equal(t, testMessage, response, "Echo response should match sent message") t.Logf("Sent: %s, Received: %s", testMessage, response) } // TestE2EWebSocketBinary 测试二进制消息。 // // 验证二进制消息正确传递。 func TestE2EWebSocketBinary(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送二进制数据 testData := []byte{0x01, 0x02, 0x03, 0x04, 0x05} err = client.SendBinary(testData) require.NoError(t, err, "Failed to send binary data") // 接收响应 response, err := client.ReceiveBinary() require.NoError(t, err, "Failed to receive binary data") assert.Equal(t, testData, response, "Echo response should match sent data") } // TestE2EWebSocketJSON 测试 JSON 消息。 // // 验证 JSON 消息正确传递。 func TestE2EWebSocketJSON(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送 JSON 数据 testData := map[string]string{"message": "hello", "id": "123"} err = client.SendJSON(testData) require.NoError(t, err, "Failed to send JSON") // 接收响应 var response map[string]string err = client.ReceiveJSON(&response) require.NoError(t, err, "Failed to receive JSON") assert.Equal(t, testData, response, "Echo response should match sent JSON") } // TestE2EWebSocketConcurrent 测试并发 WebSocket 连接。 // // 验证多个并发连接正常工作。 func TestE2EWebSocketConcurrent(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] // 创建连接池 pool, err := testutil.NewWSPool(context.Background(), wsURL, 5) require.NoError(t, err, "Failed to create WebSocket pool") defer pool.Close() // 向所有连接发送消息 testMessage := "Concurrent test" err = pool.SendAll(testMessage) require.NoError(t, err, "Failed to send to all connections") // 从所有连接接收消息 messages, err := pool.ReceiveAll() require.NoError(t, err, "Failed to receive from all connections") // 验证所有连接都收到了响应 assert.Equal(t, 5, len(messages), "Should have 5 responses") for i, msg := range messages { if msg != "" { assert.Equal(t, testMessage, msg, "Connection %d response should match", i) } } t.Logf("Received %d messages from concurrent connections", len(messages)) } // TestE2EWebSocketMultipleMessages 测试多消息传递。 // // 验证连续发送多条消息正常工作。 func TestE2EWebSocketMultipleMessages(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送多条消息 for i := 0; i < 10; i++ { msg := fmt.Sprintf("Message %d", i) err = client.Send(msg) require.NoError(t, err, "Failed to send message %d", i) response, err := client.Receive() require.NoError(t, err, "Failed to receive message %d", i) assert.Equal(t, msg, response, "Response %d should match", i) } t.Log("Successfully sent and received 10 messages") } // TestE2EWebSocketClose 测试连接关闭。 // // 验证连接正确关闭。 func TestE2EWebSocketClose(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") // 发送一条消息验证连接正常 err = client.Send("Test before close") require.NoError(t, err, "Failed to send before close") response, err := client.Receive() require.NoError(t, err, "Failed to receive before close") assert.Equal(t, "Test before close", response) // 关闭连接 err = client.Close() require.NoError(t, err, "Failed to close connection") // 验证连接已关闭 assert.True(t, client.IsClosed(), "Connection should be closed") // 尝试发送应该失败 err = client.Send("After close") assert.Error(t, err, "Send after close should fail") } // TestE2EWebSocketTimeout 测试超时处理。 // // 验证超时配置生效。 func TestE2EWebSocketTimeout(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送消息 err = client.Send("Timeout test") require.NoError(t, err, "Failed to send message") // 使用短超时接收 response, err := client.ReceiveWithTimeout(5 * time.Second) require.NoError(t, err, "Failed to receive with timeout") assert.Equal(t, "Timeout test", response) } // TestE2EWebSocketHeaders 测试自定义头部。 // // 验证自定义请求头正确传递。 func TestE2EWebSocketHeaders(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] // 创建带自定义头部的客户端 headers := http.Header{} headers.Set("X-Custom-Header", "test-value") headers.Set("Authorization", "Bearer token123") client, err := testutil.NewWSClient(context.Background(), wsURL, testutil.WithWSHeaders(headers), ) require.NoError(t, err, "Failed to create WebSocket client with headers") defer client.Close() // 发送消息验证连接正常 err = client.Send("Headers test") require.NoError(t, err, "Failed to send message") response, err := client.Receive() require.NoError(t, err, "Failed to receive message") assert.Equal(t, "Headers test", response) } // TestE2EWebSocketPoolOperations 测试连接池操作。 // // 验证连接池的各种操作。 func TestE2EWebSocketPoolOperations(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] pool, err := testutil.NewWSPool(context.Background(), wsURL, 3) require.NoError(t, err, "Failed to create WebSocket pool") defer pool.Close() // 验证连接数量 assert.Equal(t, 3, pool.Count(), "Pool should have 3 connections") // 向单个连接发送消息 err = pool.SendOne(0, "Single message") require.NoError(t, err, "Failed to send to single connection") // 向所有连接发送消息 err = pool.SendAll("Broadcast message") require.NoError(t, err, "Failed to broadcast") } // TestE2EWebSocketReconnect 测试重连场景。 // // 验证客户端可以重新建立连接。 func TestE2EWebSocketReconnect(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] // 第一次连接 client1, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create first client") err = client1.Send("First connection") require.NoError(t, err, "Failed to send on first connection") response, err := client1.Receive() require.NoError(t, err, "Failed to receive on first connection") assert.Equal(t, "First connection", response) // 关闭第一个连接 client1.Close() // 重新建立连接 client2, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create second client") defer client2.Close() err = client2.Send("Second connection") require.NoError(t, err, "Failed to send on second connection") response, err = client2.Receive() require.NoError(t, err, "Failed to receive on second connection") assert.Equal(t, "Second connection", response) t.Log("Reconnect test completed successfully") } // TestE2EWebSocketLargeMessage 测试大消息。 // // 验证大消息正确传递。 func TestE2EWebSocketLargeMessage(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(wsEchoHandler)) defer server.Close() wsURL := "ws" + server.URL[4:] client, err := testutil.NewWSClient(context.Background(), wsURL) require.NoError(t, err, "Failed to create WebSocket client") defer client.Close() // 发送大消息(64KB) largeData := make([]byte, 64*1024) for i := range largeData { largeData[i] = byte(i % 256) } err = client.SendBinary(largeData) require.NoError(t, err, "Failed to send large data") response, err := client.ReceiveBinary() require.NoError(t, err, "Failed to receive large data") assert.Equal(t, len(largeData), len(response), "Response size should match") assert.Equal(t, largeData, response, "Response data should match") } // TestE2EWebSocketProxyIntegration 测试 WebSocket 代理集成。 // // 注意:此测试需要 Docker 环境,验证 lolly WebSocket 代理功能。 func TestE2EWebSocketProxyIntegration(t *testing.T) { t.Parallel() if testing.Short() { t.Skip("Skipping E2E WebSocket proxy test in short mode") } ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second) defer cancel() if !testutil.LollyImageAvailable(ctx) { t.Skip("lolly:latest image not available, run 'make docker-build' first") } // 启动后端 netObj, networkName, pool, err := testutil.SetupProxyTest(ctx, 1, t.Name()) require.NoError(t, err, "Failed to start backend pool") defer testutil.CleanupProxyTest(ctx, netObj, networkName, pool) // 构建配置 cfg := testutil.NewConfigBuilder(). WithServer(":8080"). WithProxy("/", pool.InternalAddresses()) configYAML, err := cfg.Build() require.NoError(t, err, "Failed to build config") // 启动 lolly lolly, err := testutil.StartLolly(ctx, testutil.WithConfigYAML(configYAML), testutil.WithNetwork(networkName)) require.NoError(t, err, "Failed to start lolly") defer lolly.Terminate(ctx) err = lolly.WaitForHealthy(ctx, testutil.HealthCheckWaitTimeout) require.NoError(t, err, "Lolly not healthy") // 测试 HTTP 代理(WebSocket 需要 WebSocket 后端) client := &http.Client{Timeout: 10 * time.Second} resp, err := client.Get(lolly.HTTPBaseURL()) require.NoError(t, err, "HTTP request failed") defer resp.Body.Close() assert.Equal(t, 200, resp.StatusCode) t.Log("WebSocket proxy integration test placeholder - requires WebSocket backend") }