Agent-System/internal/squad/workflow_test.go

package squad_test

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/squad"
)

func TestWorkflowDriver_DynamicAndExecute(t *testing.T) {
	tmpDir := t.TempDir()

	engine := squad.NewWorkspaceEngine(nil)
	ws, err := engine.PrepareWorkspace(squad.WorkspacePrepareRequest{
		WorkDir:      tmpDir,
		TaskKey:      "intake",
		StageTitle:   "用户初始需求收集表",
		ProjectName:  "无创神经刺激器",
		CustomerName: "圆山研究院",
		Backend:      "stub",
		WorkflowSteps: []squad.StepSpec{
			{
				ID:           "step_01_organize",
				Role:         "organizer",
				Title:        "抽取事实并回填字段",
				SubWorkspace: "tasks/01_organize",
				Mode:         "execute",
			},
			{
				ID:           "step_02_draft",
				Role:         "writer",
				Title:        "撰写初稿",
				SubWorkspace: "tasks/02_draft",
				Mode:         "execute",
				Needs:        []string{"tasks/01_organize/output.md"},
			},
		},
	})
	if err != nil {
		t.Fatalf("PrepareWorkspace failed: %v", err)
	}

	driver := squad.NewWorkflowDriver(engine)

	// 1. 动态增加一个审核步骤
	err = driver.AddStep(ws.WorkDir, squad.StepSpec{
		ID:           "step_03_review",
		Role:         "leader",
		Title:        "主管终审",
		SubWorkspace: "tasks/03_review",
		Mode:         "execute",
		Needs:        []string{"tasks/02_draft/output.md"},
	})
	if err != nil {
		t.Fatalf("AddStep failed: %v", err)
	}

	// 校验 workflow.json 是否更新为 3 步
	wf, err := driver.LoadWorkflow(ws.WorkDir)
	if err != nil {
		t.Fatalf("LoadWorkflow failed: %v", err)
	}
	if len(wf.Steps) != 3 {
		t.Fatalf("expected 3 steps after AddStep, got %d", len(wf.Steps))
	}

	// 模拟 organizer 在 tasks/01_organize 下提前预置 fields.json (验证黑板自动捕获)
	task1Dir := filepath.Join(ws.WorkDir, "tasks", "01_organize")
	_ = os.MkdirAll(task1Dir, 0o755)
	fieldsData := map[string]string{
		"customer_name": "圆山研究院",
		"device_type":   "二类医疗器械",
		"intended_use":  "改善打鼾与夜间睡眠",
	}
	fBytes, _ := json.Marshal(fieldsData)
	_ = os.WriteFile(filepath.Join(task1Dir, "fields.json"), fBytes, 0o644)

	// 2. 执行整个工作流
	stubBackend := &agent.StubBackend{}
	var stepEvents []squad.StepEvent
	opts := squad.ExecuteOptions{
		Timeout: 5 * time.Second,
		StepNotify: func(ev squad.StepEvent) {
			stepEvents = append(stepEvents, ev)
		},
	}

	res, err := driver.Execute(context.Background(), ws.WorkDir, stubBackend, opts)
	if err != nil {
		t.Fatalf("driver.Execute failed: %v", err)
	}

	if !res.Success {
		t.Fatalf("expected execution success, error: %s", res.ErrorReason)
	}
	if res.StepsRun != 3 {
		t.Errorf("expected 3 steps run, got %d", res.StepsRun)
	}

	// 3. 校验每个子工作区是否生成了 output.md
	for _, sub := range []string{"01_organize", "02_draft", "03_review"} {
		p := filepath.Join(ws.WorkDir, "tasks", sub, "output.md")
		if _, err := os.Stat(p); err != nil {
			t.Errorf("expected %s/output.md to exist: %v", sub, err)
		}
	}

	// 4. 校验 blackboard.json 是否成功回填了 fields.json 中的字段
	bb, err := driver.LoadBlackboard(ws.WorkDir)
	if err != nil {
		t.Fatalf("LoadBlackboard failed: %v", err)
	}
	if bb.Fields["device_type"] != "二类医疗器械" {
		t.Errorf("expected blackboard device_type '二类医疗器械', got '%s'", bb.Fields["device_type"])
	}

	// 5. 校验根目录下是否汇总生成了 output.md
	rootOut := filepath.Join(ws.WorkDir, "output.md")
	if _, err := os.Stat(rootOut); err != nil {
		t.Fatalf("root output.md does not exist: %v", err)
	}
}