1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  //
  3  // Licensed under the Apache License, Version 2.0 (the "License");
  4  // you may not use this file except in compliance with the License.
  5  // You may obtain a copy of the License at
  6  //
  7  //     http://www.apache.org/licenses/LICENSE-2.0
  8  //
  9  // Unless required by applicable law or agreed to in writing, software
 10  // distributed under the License is distributed on an "AS IS" BASIS,
 11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  // See the License for the specific language governing permissions and
 13  // limitations under the License.
 14  //
 15  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  package agent
 20  
 21  import (
 22  	"encoding/json"
 23  	"fmt"
 24  	"strings"
 25  	"time"
 26  
 27  	"testing"
 28  
 29  	"github.com/stretchr/testify/assert"
 30  )
 31  
 32  // TestLargeDataSend 测试发送大字节数据
 33  func TestLargeDataSend(t *testing.T) {
 34  	// 创建Agent实例（不连接到真实服务器）
 35  	agent := NewAgent(AgentConfig{
 36  		ServerURL: "ws://xx/api/v1/agents/ws", // 使用测试URL
 37  		Info: AgentInfo{
 38  			ID:       "test-large-data",
 39  			HostName: "test-host",
 40  			IP:       "127.0.0.1",
 41  			Version:  "0.1",
 42  			Metadata: "",
 43  		},
 44  	})
 45  	err := agent.connect()
 46  	assert.NoError(t, err)
 47  	// 启动各种协程
 48  	go agent.handleSend()
 49  	go agent.handleReceive()
 50  
 51  	// 创建大数据内容 - 生成约1MB的数据
 52  	largeContent := generateLargeContent(1024 * 1024) // 1MB
 53  
 54  	// 创建包含大数据的任务结果
 55  	largeResult := map[string]interface{}{
 56  		"type":        "large_data_test",
 57  		"timestamp":   time.Now().Unix(),
 58  		"data_size":   len(largeContent),
 59  		"content":     largeContent,
 60  		"description": "测试发送大字节数据的能力",
 61  		"metadata": map[string]interface{}{
 62  			"compression": false,
 63  			"encoding":    "utf-8",
 64  			"chunks":      1,
 65  		},
 66  	}
 67  
 68  	// 测试序列化大数据
 69  	jsonData, err := json.Marshal(largeResult)
 70  	assert.NoError(t, err, "大数据JSON序列化应该成功")
 71  
 72  	dataSize := len(jsonData)
 73  	t.Logf("生成的JSON数据大小: %d bytes (%.2f MB)", dataSize, float64(dataSize)/(1024*1024))
 74  
 75  	// 测试通过sendChan发送大数据（模拟真实发送）
 76  	sessionId := "test-session-large-data"
 77  
 78  	// 发送大数据
 79  	err = agent.SendTaskResult(sessionId, largeResult)
 80  	assert.NoError(t, err, "发送大数据任务结果应该成功")
 81  
 82  	// 等待一小段时间确保消息被处理
 83  	time.Sleep(5 * time.Second)
 84  
 85  	t.Log("大字节数据发送测试完成")
 86  }
 87  
 88  // generateLargeContent 生成指定大小的大内容
 89  func generateLargeContent(size int) string {
 90  	// 创建基础模板
 91  	template := "这是一段测试数据，用于验证大字节数据的传输能力。包含中文字符以测试编码处理。Data chunk %d. "
 92  
 93  	var builder strings.Builder
 94  	builder.Grow(size) // 预分配容量
 95  
 96  	chunkCount := 0
 97  	for builder.Len() < size {
 98  		chunk := fmt.Sprintf(template, chunkCount)
 99  		if builder.Len()+len(chunk) > size {
100  			// 添加剩余字符直到达到目标大小
101  			remaining := size - builder.Len()
102  			builder.WriteString(chunk[:remaining])
103  			break
104  		}
105  		builder.WriteString(chunk)
106  		chunkCount++
107  	}
108  
109  	return builder.String()
110  }