Agent-System/internal/stream/stream_test.go

package stream

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"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_EmitError_PopulatesErrorInfo(t *testing.T) {
	rawErr := errors.New("invalid project config")
	var received StreamEvent

	streamer := NewStreamer(true, func(ev StreamEvent) {
		received = ev
	})

	streamer.EmitError(rawErr)

	if received.Type != EventError {
		t.Fatalf("type=%s", received.Type)
	}
	if received.ErrorInfo == nil {
		t.Fatal("ErrorInfo is nil")
	}
	if received.Error != received.ErrorInfo.UserMessage {
		t.Fatalf("Error=%q UserMessage=%q", received.Error, received.ErrorInfo.UserMessage)
	}
	if received.ErrorInfo.Upstream != rawErr.Error() {
		t.Fatalf("Upstream=%q want %q", received.ErrorInfo.Upstream, rawErr.Error())
	}
	if received.ErrorInfo.Class == "" {
		t.Fatal("Class is empty")
	}
}

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)
	}
}

func TestStreamer_EmitDoneWithWorkspace(t *testing.T) {
	var received StreamEvent

	streamer := NewStreamer(true, func(ev StreamEvent) {
		received = ev
	})

	streamer.EmitDoneWithWorkspace("task_run", "/tmp/ws-demo", 1200*time.Millisecond)

	if received.Type != EventDone {
		t.Fatalf("type=%s", received.Type)
	}
	if received.Done == nil {
		t.Fatal("Done payload is nil")
	}
	if received.Done.WorkspaceID != "/tmp/ws-demo" {
		t.Fatalf("WorkspaceID=%q want %q", received.Done.WorkspaceID, "/tmp/ws-demo")
	}
	if received.Done.DurationMs != 1200 {
		t.Fatalf("DurationMs=%d want 1200", received.Done.DurationMs)
	}

	// workspace_id 为空时 omitempty 不应出现在 JSON 中
	streamer.EmitDone("task_run", 100*time.Millisecond)
	data, err := json.Marshal(received)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	if strings.Contains(string(data), "workspace_id") {
		t.Fatalf("expected workspace_id omitted, got %s", data)
	}
}