Files
Agent-System/internal/storage/cas_test.go
T

285 lines
8.6 KiB
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)
}
}
// TestCAS_ImmutabilityCoWContract 固化 CoW(写时复制)契约:
// 1. 同内容重复写入只增引用计数,不写第二份物理文件,且旧 blob 字节保持不变;
// 2. 可编辑文档产生新内容时按新 hash 落新 blob(写入即复制),旧版本 blob
// 不被覆写,历史引用仍可读出原始字节。
//
// 该测试是防回归守卫:未来任何「就地改写 blob 文件」的实现都会在此失败。
func TestCAS_ImmutabilityCoWContract(t *testing.T) {
tempDir, err := os.MkdirTemp("", "cas_cow_test_*")
if err != nil {
t.Fatalf("mkdir temp failed: %v", err)
}
defer os.RemoveAll(tempDir)
cas, err := NewCAS(tempDir, nil, nil)
if err != nil {
t.Fatalf("new cas failed: %v", err)
}
ctx := context.Background()
v1 := []byte("V1:产品需求说明书初始版本内容")
refA, err := cas.StoreFile(ctx, 900, 1, "prd.docx", "material", v1, "application/octet-stream")
if err != nil {
t.Fatalf("store v1 failed: %v", err)
}
blobPath := filepath.Join(tempDir, "blobs", refA.BlobSHA256[:2], refA.BlobSHA256[2:4], refA.BlobSHA256+".bin")
v1Bytes, err := os.ReadFile(blobPath)
if err != nil {
t.Fatalf("v1 blob file missing: %v", err)
}
// 1. 相同内容再传:命中去重,物理文件数量与字节内容不变。
refB, err := cas.StoreFile(ctx, 900, 2, "prd_copy.docx", "material", v1, "application/octet-stream")
if err != nil {
t.Fatalf("store duplicate failed: %v", err)
}
if !refB.IsDeduplicated || refB.BlobSHA256 != refA.BlobSHA256 {
t.Fatalf("expected dedup hit on identical content")
}
afterDup, err := os.ReadFile(blobPath)
if err != nil || string(afterDup) != string(v1Bytes) {
t.Fatalf("duplicate upload mutated existing blob bytes (CoW violation)")
}
// 2. 编辑后写入(内容变化):必须产生新 blob;旧 blob 不被覆写。
v2 := []byte("V2:产品需求说明书——评审后修改条款")
refC, err := cas.StoreFile(ctx, 900, 1, "prd.docx", "material", v2, "application/octet-stream")
if err != nil {
t.Fatalf("store v2 failed: %v", err)
}
if refC.BlobSHA256 == refA.BlobSHA256 {
t.Fatalf("modified content must get a new blob hash")
}
if refC.IsDeduplicated {
t.Fatalf("modified content must not be marked dedup")
}
stillV1, err := os.ReadFile(blobPath)
if err != nil || string(stillV1) != string(v1Bytes) {
t.Fatalf("write of new version overwrote old blob (CoW violation)")
}
gotV1, _, err := cas.ReadRefBytes(ctx, refA.ID)
if err != nil || string(gotV1) != string(v1) {
t.Fatalf("historical ref must still read original V1 bytes, got %q err=%v", gotV1, err)
}
gotV2, _, err := cas.ReadRefBytes(ctx, refC.ID)
if err != nil || string(gotV2) != string(v2) {
t.Fatalf("new ref must read V2 bytes, got %q err=%v", gotV2, err)
}
// 3. 物理盘上恰好两份 blob(V1 一份共享 + V2 一份)。
count := 0
_ = filepath.Walk(filepath.Join(tempDir, "blobs"), func(path string, info os.FileInfo, err error) error {
if err == nil && !info.IsDir() && filepath.Ext(path) == ".bin" {
count++
}
return nil
})
if count != 2 {
t.Fatalf("expected exactly 2 physical blobs (CoW), got %d", count)
}
}