Agent-System/internal/storage/cas_test.go

package storage

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"testing"
	"time"
)

type mockEmbedder struct {
	calledCount int
}

func (m *mockEmbedder) EmbedOne(ctx context.Context, text string) ([]float32, error) {
	m.calledCount++
	// 返回一个简单的测试向量
	vec := make([]float32, 1024)
	for i := range vec {
		vec[i] = 0.01 * float32(i%10)
	}
	return vec, nil
}

func TestCAS_FileDeduplication(t *testing.T) {
	tempDir, err := os.MkdirTemp("", "cas_test_*")
	if err != nil {
		t.Fatalf("mkdir temp failed: %v", err)
	}
	defer os.RemoveAll(tempDir)

	mockEmb := &mockEmbedder{}
	cas, err := NewCAS(tempDir, nil, mockEmb)
	if err != nil {
		t.Fatalf("new cas failed: %v", err)
	}

	ctx := context.Background()
	fileContent := []byte("这是一份用于去重测试的医疗器械说明书与注册申报资料内容，包含重复章节条款。")

	// 1. 用户 1 首次上传
	ref1, err := cas.StoreFile(ctx, 101, 1, "instruction_v1.pdf", "material", fileContent, "application/pdf")
	if err != nil {
		t.Fatalf("store file 1 failed: %v", err)
	}
	if ref1.IsDeduplicated {
		t.Errorf("expected first upload to not be deduplicated")
	}

	// 2. 用户 2 再次上传完全相同内容的文件 (不同文件名、不同用户)
	ref2, err := cas.StoreFile(ctx, 101, 2, "instruction_copy.pdf", "material", fileContent, "application/pdf")
	if err != nil {
		t.Fatalf("store file 2 failed: %v", err)
	}
	if !ref2.IsDeduplicated {
		t.Errorf("expected second upload to be deduplicated (CAS Hit)")
	}
	if ref1.BlobSHA256 != ref2.BlobSHA256 {
		t.Errorf("expected identical SHA256 hashes, got %s vs %s", ref1.BlobSHA256, ref2.BlobSHA256)
	}

	// 3. 读取验证文件完整性
	r1, _, err := cas.OpenFileByRef(ctx, ref1.ID)
	if err != nil {
		t.Fatalf("open file 1 failed: %v", err)
	}
	defer r1.Close()
	data1, _ := io.ReadAll(r1)
	if string(data1) != string(fileContent) {
		t.Errorf("file 1 content mismatch")
	}

	r2, _, err := cas.OpenFileByRef(ctx, ref2.ID)
	if err != nil {
		t.Fatalf("open file 2 failed: %v", err)
	}
	defer r2.Close()
	data2, _ := io.ReadAll(r2)
	if string(data2) != string(fileContent) {
		t.Errorf("file 2 content mismatch")
	}

	// 4. 检查物理磁盘上是否仅有 1 个物理 blob 文件
	blobsDir := filepath.Join(tempDir, "blobs")
	blobCount := 0
	_ = filepath.Walk(blobsDir, func(path string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() && filepath.Ext(path) == ".bin" {
			blobCount++
		}
		return nil
	})
	if blobCount != 1 {
		t.Errorf("expected exactly 1 physical blob on disk, found %d", blobCount)
	}
}

func TestCAS_MemoryChunkAndDecay(t *testing.T) {
	tempDir, err := os.MkdirTemp("", "cas_chunk_test_*")
	if err != nil {
		t.Fatalf("mkdir temp failed: %v", err)
	}
	defer os.RemoveAll(tempDir)

	mockEmb := &mockEmbedder{}
	cas, err := NewCAS(tempDir, nil, mockEmb)
	if err != nil {
		t.Fatalf("new cas failed: %v", err)
	}

	ctx := context.Background()

	// 1. 存储切片
	c1, err := cas.StoreMemoryChunk(ctx, MemoryChunkInput{
		Domain:           "user_memory",
		ProjectID:        1001,
		UserID:           1,
		Title:            "客户对工期极度敏感",
		Content:          "客户希望首期在 4 周内完成样品摸底，并强调需优先保证安规测试通过。",
		ImportanceWeight: 1.5,
	})
	if err != nil {
		t.Fatalf("store chunk 1 failed: %v", err)
	}
	if c1.ChunkHash == "" {
		t.Errorf("expected non-empty chunk hash")
	}

	// 2. 模拟检索
	hits, err := cas.SearchMemory(ctx, MemorySearchOptions{
		Domain:            "user_memory",
		ProjectID:         1001,
		QueryText:         "工期",
		ApplyTimeDecay:    true,
		DecayHalfLifeDays: 7.0,
		TopK:              5,
	})
	if err != nil {
		t.Fatalf("search memory failed: %v", err)
	}
	if len(hits) == 0 {
		t.Fatalf("expected memory hits, got 0")
	}
	if hits[0].Chunk.Title != "客户对工期极度敏感" {
		t.Errorf("unexpected top hit title: %s", hits[0].Chunk.Title)
	}
	if hits[0].FinalScore <= 0 {
		t.Errorf("expected positive final score")
	}
}

func TestCAS_TimeDecayCalculation(t *testing.T) {
	tempDir, err := os.MkdirTemp("", "cas_decay_test_*")
	if err != nil {
		t.Fatalf("mkdir temp failed: %v", err)
	}
	defer os.RemoveAll(tempDir)

	cas, _ := NewCAS(tempDir, nil, nil)
	ctx := context.Background()

	// 写入一条即时记忆和一条较早的记忆
	nowChunk, _ := cas.StoreMemoryChunk(ctx, MemoryChunkInput{
		Domain:           "project_activity",
		ProjectID:        2001,
		Title:            "今日决策",
		Content:          "今日锁定目标市场为中国及欧盟CE。",
		ImportanceWeight: 1.0,
	})

	pastChunk, _ := cas.StoreMemoryChunk(ctx, MemoryChunkInput{
		Domain:           "project_activity",
		ProjectID:        2001,
		Title:            "早期设想",
		Content:          "早期曾考虑走美国FDA 510(k)途径，后续搁置。",
		ImportanceWeight: 1.0,
	})

	// 人工修改 pastChunk 的创建时间为 14 天前 (两个半衰期)
	pastChunk.CreatedAt = time.Now().Add(-14 * 24 * time.Hour)

	hits, err := cas.SearchMemory(ctx, MemorySearchOptions{
		Domain:            "project_activity",
		ProjectID:         2001,
		ApplyTimeDecay:    true,
		DecayHalfLifeDays: 7.0,
		TopK:              5,
	})
	if err != nil {
		t.Fatalf("search memory failed: %v", err)
	}
	if len(hits) < 2 {
		t.Fatalf("expected 2 hits, got %d", len(hits))
	}

	// 即时记忆得分应远高于 14 天前的早期设想
	if hits[0].Chunk.ID != nowChunk.ID {
		t.Errorf("expected nowChunk to rank higher due to time decay, got ID %d", hits[0].Chunk.ID)
	}
	if hits[1].TimeDecay >= 0.5 {
		t.Errorf("expected 14-day old chunk to decay below 0.5, got %.3f", hits[1].TimeDecay)
	}
}