agent-team/internal/agent/agent_test.go

package agent

import (
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestMemory_SmartLoad(t *testing.T) {
	// 创建临时 agent 目录
	dir := t.TempDir()
	memDir := filepath.Join(dir, "memory")
	os.MkdirAll(memDir, 0755)

	a := &Agent{Dir: dir}

	// 写入 core.md
	os.WriteFile(filepath.Join(memDir, "core.md"), []byte("核心经验1\n"), 0644)

	// 写入当月文件
	currentMonth := time.Now().Format("2006-01") + ".md"
	os.WriteFile(filepath.Join(memDir, currentMonth), []byte("当月记忆\n"), 0644)

	// 写入旧月份文件（不应被加载）
	os.WriteFile(filepath.Join(memDir, "2025-01.md"), []byte("旧记忆不应出现\n"), 0644)

	mem := a.Memory()

	if !strings.Contains(mem, "核心经验1") {
		t.Error("应包含 core.md 内容")
	}
	if !strings.Contains(mem, "当月记忆") {
		t.Error("应包含当月文件内容")
	}
	if strings.Contains(mem, "旧记忆不应出现") {
		t.Error("不应包含旧月份文件内容")
	}
}

func TestMemory_HardLimit(t *testing.T) {
	dir := t.TempDir()
	memDir := filepath.Join(dir, "memory")
	os.MkdirAll(memDir, 0755)

	a := &Agent{Dir: dir}

	// 写入超大 core.md（5000 个中文字符）
	bigContent := strings.Repeat("记", 5000)
	os.WriteFile(filepath.Join(memDir, "core.md"), []byte(bigContent), 0644)

	mem := a.Memory()
	runes := []rune(mem)

	// 应被截断到 4000 + 截断提示
	if len(runes) > memoryMaxChars+50 {
		t.Errorf("记忆应被截断，实际 %d 字符", len(runes))
	}
	if !strings.Contains(mem, "记忆已截断") {
		t.Error("应包含截断提示")
	}
}

func TestMemory_EmptyDir(t *testing.T) {
	dir := t.TempDir()
	a := &Agent{Dir: dir}

	// memory 目录不存在
	mem := a.Memory()
	if mem != "" {
		t.Error("不存在 memory 目录时应返回空字符串")
	}
}

func TestKnowledge_Load(t *testing.T) {
	dir := t.TempDir()
	knowledgeDir := filepath.Join(dir, "knowledge")
	os.MkdirAll(knowledgeDir, 0755)

	os.WriteFile(filepath.Join(knowledgeDir, "法律基础.md"), []byte("合同法要点\n"), 0644)
	os.WriteFile(filepath.Join(knowledgeDir, "案例.md"), []byte("经典案例\n"), 0644)
	os.WriteFile(filepath.Join(knowledgeDir, "ignore.txt"), []byte("不是md\n"), 0644)

	a := &Agent{Dir: dir, KnowledgeDir: knowledgeDir}

	k := a.Knowledge()

	if !strings.Contains(k, "合同法要点") {
		t.Error("应包含法律基础.md 内容")
	}
	if !strings.Contains(k, "经典案例") {
		t.Error("应包含案例.md 内容")
	}
	if strings.Contains(k, "不是md") {
		t.Error("不应包含非 .md 文件")
	}
}

func TestKnowledge_Empty(t *testing.T) {
	a := &Agent{Dir: t.TempDir()}

	// KnowledgeDir 未设置
	if a.Knowledge() != "" {
		t.Error("KnowledgeDir 为空时应返回空字符串")
	}
}

func TestKnowledge_HardLimit(t *testing.T) {
	dir := t.TempDir()
	knowledgeDir := filepath.Join(dir, "knowledge")
	os.MkdirAll(knowledgeDir, 0755)

	bigContent := strings.Repeat("知", 10000)
	os.WriteFile(filepath.Join(knowledgeDir, "big.md"), []byte(bigContent), 0644)

	a := &Agent{Dir: dir, KnowledgeDir: knowledgeDir}
	k := a.Knowledge()
	runes := []rune(k)

	if len(runes) > knowledgeMaxChars+50 {
		t.Errorf("知识库应被截断，实际 %d 字符", len(runes))
	}
	if !strings.Contains(k, "知识库已截断") {
		t.Error("应包含截断提示")
	}
}

func TestBuildSystemPrompt_WithKnowledgeAndMemory(t *testing.T) {
	dir := t.TempDir()
	memDir := filepath.Join(dir, "memory")
	knowledgeDir := filepath.Join(dir, "knowledge")
	os.MkdirAll(memDir, 0755)
	os.MkdirAll(knowledgeDir, 0755)

	// 写入 SOUL.md
	os.WriteFile(filepath.Join(dir, "SOUL.md"), []byte("你是测试助手。"), 0644)

	// 写入 knowledge
	os.WriteFile(filepath.Join(knowledgeDir, "test.md"), []byte("测试知识"), 0644)

	// 写入当月 memory
	currentMonth := time.Now().Format("2006-01") + ".md"
	os.WriteFile(filepath.Join(memDir, currentMonth), []byte("测试记忆"), 0644)

	a := &Agent{Dir: dir, KnowledgeDir: knowledgeDir, Soul: "你是测试助手。"}
	prompt := a.BuildSystemPrompt("额外上下文")

	// 验证顺序：soul → knowledge → memory → extra
	knowledgeIdx := strings.Index(prompt, "<knowledge>")
	memoryIdx := strings.Index(prompt, "<memory>")
	extraIdx := strings.Index(prompt, "额外上下文")

	if knowledgeIdx < 0 {
		t.Fatal("prompt 中应包含 <knowledge> 标签")
	}
	if memoryIdx < 0 {
		t.Fatal("prompt 中应包含 <memory> 标签")
	}
	if knowledgeIdx >= memoryIdx {
		t.Error("<knowledge> 应在 <memory> 之前")
	}
	if memoryIdx >= extraIdx {
		t.Error("<memory> 应在额外上下文之前")
	}
	if !strings.Contains(prompt, "测试知识") {
		t.Error("prompt 应包含知识库内容")
	}
	if !strings.Contains(prompt, "测试记忆") {
		t.Error("prompt 应包含记忆内容")
	}
}