141 lines
3.7 KiB
Go
141 lines
3.7 KiB
Go
package stream
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestStreamer_BasicEvents(t *testing.T) {
|
|
var received []StreamEvent
|
|
var mu sync.Mutex
|
|
|
|
listener := func(ev StreamEvent) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
received = append(received, ev)
|
|
}
|
|
|
|
streamer := NewStreamer(true, listener)
|
|
|
|
streamer.EmitPhase("squad_assembly", "正在组建小队...", "leader", 1, 3, 33)
|
|
streamer.EmitStatus("ADVANCE_STAGE", "BranchExpectationHit", "V0.5_feasibility", "推进阶段", map[string]string{"market": "CE"})
|
|
streamer.EmitSummary("初审结论", "产品符合第二类医疗器械判定", 80, "V0.5")
|
|
streamer.EmitBriefing(map[string]any{"maturity_score": 80, "maturity_level": "V0.5"})
|
|
streamer.EmitDetail("step_output", "tasks/01_regulatory", "regulatory", "条款梳理", "GB 9706.1 适用", true)
|
|
streamer.EmitDelta("您好，已为您完成推进。")
|
|
streamer.EmitDeliverable("feasibility_reg", "注册可行性评估报告", "sandbox/output.md", 2500, 100.0)
|
|
streamer.EmitError(errors.New("模拟告警"))
|
|
streamer.EmitDone("sess_123", 500*time.Millisecond)
|
|
|
|
mu.Lock()
|
|
count := len(received)
|
|
mu.Unlock()
|
|
|
|
if count != 9 {
|
|
t.Fatalf("expected 9 events, got %d", count)
|
|
}
|
|
|
|
// 校验事件类型
|
|
expectedTypes := []EventType{
|
|
EventPhase, EventStatus, EventSummary, EventBriefing,
|
|
EventDetail, EventDelta, EventDeliverable, EventError, EventDone,
|
|
}
|
|
|
|
for i, exp := range expectedTypes {
|
|
if received[i].Type != exp {
|
|
t.Errorf("event %d: expected type %s, got %s", i, exp, received[i].Type)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestStreamer_FilterDetails(t *testing.T) {
|
|
var received []StreamEvent
|
|
|
|
listener := func(ev StreamEvent) {
|
|
received = append(received, ev)
|
|
}
|
|
|
|
// includeDetails = false: 必须过滤 Detail 事件
|
|
streamer := NewStreamer(false, listener)
|
|
|
|
streamer.EmitPhase("step_executing", "法规专家研判中...", "regulatory", 1, 2, 50)
|
|
streamer.EmitDetail("step_output", "tasks/01", "regulatory", "庞大推演草稿", "长篇草稿内容...", true)
|
|
streamer.EmitSummary("精炼结论", "完成研判", 85, "V0.5")
|
|
|
|
if len(received) != 2 {
|
|
t.Fatalf("expected 2 events (Detail filtered out), got %d", len(received))
|
|
}
|
|
|
|
if received[0].Type != EventPhase || received[1].Type != EventSummary {
|
|
t.Errorf("unexpected event types: %v, %v", received[0].Type, received[1].Type)
|
|
}
|
|
}
|
|
|
|
func TestStreamer_ContextPropagation(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// 1. 未注入时 FromContext 返回 nil，测试 nil-safe 调用
|
|
nilStreamer := FromContext(ctx)
|
|
if nilStreamer != nil {
|
|
t.Fatalf("expected nil streamer from empty context")
|
|
}
|
|
// nil 调用不崩溃
|
|
nilStreamer.EmitPhase("test", "msg", "role", 1, 1, 100)
|
|
nilStreamer.EmitDelta("test")
|
|
|
|
// 2. 注入后正常提取
|
|
var events []StreamEvent
|
|
ctxWithS, s := WithListener(ctx, true, func(ev StreamEvent) {
|
|
events = append(events, ev)
|
|
})
|
|
|
|
extracted := FromContext(ctxWithS)
|
|
if extracted != s {
|
|
t.Fatalf("extracted streamer did not match original")
|
|
}
|
|
|
|
extracted.EmitPhase("squad_assembly", "组建小队", "leader", 1, 3, 33)
|
|
if len(events) != 1 {
|
|
t.Fatalf("expected 1 event, got %d", len(events))
|
|
}
|
|
}
|
|
|
|
func TestStreamer_Concurrency(t *testing.T) {
|
|
var count int64
|
|
var mu sync.Mutex
|
|
|
|
s := NewStreamer(true, func(ev StreamEvent) {
|
|
mu.Lock()
|
|
count++
|
|
mu.Unlock()
|
|
})
|
|
|
|
var wg sync.WaitGroup
|
|
workers := 10
|
|
iterations := 50
|
|
|
|
for i := 0; i < workers; i++ {
|
|
wg.Add(1)
|
|
go func(wID int) {
|
|
defer wg.Done()
|
|
for j := 0; j < iterations; j++ {
|
|
s.EmitPhase("running", "executing step", "worker", j, iterations, 50)
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
mu.Lock()
|
|
total := count
|
|
mu.Unlock()
|
|
|
|
expected := int64(workers * iterations)
|
|
if total != expected {
|
|
t.Fatalf("expected %d events, got %d", expected, total)
|
|
}
|
|
}
|