Agent-System/internal/storage/cas_db_test.go

package storage

import (
	"context"
	"os"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

func TestCAS_DatabaseIntegration(t *testing.T) {
	dsn := os.Getenv("CONTEXT_DB_DSN")
	if dsn == "" {
		dsn = os.Getenv("TEST_DATABASE_URL")
	}
	if dsn == "" {
		t.Skip("CONTEXT_DB_DSN or TEST_DATABASE_URL must be set (refusing default dev database)")
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Skipf("cannot connect to context db (%v), skipping DB integration test", err)
		return
	}
	defer pool.Close()

	if err := pool.Ping(ctx); err != nil {
		t.Skipf("context db ping failed (%v), skipping DB integration test", err)
		return
	}

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

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

	// 准备测试项目与用户
	var projID int64
	projCode := "TEST_PROJ_" + time.Now().Format("20060102150405")
	err = pool.QueryRow(ctx, `
		INSERT INTO projects (code, name, customer_name, device_type)
		VALUES ($1, '测试内窥镜系统', '北京水木东方', '内窥镜')
		ON CONFLICT (code) DO UPDATE SET name = EXCLUDED.name
		RETURNING id
	`, projCode).Scan(&projID)
	if err != nil {
		t.Fatalf("insert project failed: %v", err)
	}

	var userID int64
	accountID := "user_test_" + time.Now().Format("150405")
	err = pool.QueryRow(ctx, `
		INSERT INTO users (account_id, name, role)
		VALUES ($1, '测试工程师', 'engineer')
		ON CONFLICT (account_id) DO UPDATE SET name = EXCLUDED.name
		RETURNING id
	`, accountID).Scan(&userID)
	if err != nil {
		t.Fatalf("insert user failed: %v", err)
	}

	// 1. 上传文件 1 (使用动态内容确保本次测试为全新文件)
	testData := []byte("这是写入真实数据库并在物理磁盘去重的测试申报文件内容: " + time.Now().Format(time.RFC3339Nano))
	ref1, err := cas.StoreFile(ctx, projID, userID, "spec_v1.pdf", "standard", testData, "application/pdf")
	if err != nil {
		t.Fatalf("store file 1 to db failed: %v", err)
	}
	if ref1.IsDeduplicated {
		t.Errorf("expected first upload to not be deduplicated")
	}

	// 2. 上传文件 2 (内容相同，触发 DB 去重)
	ref2, err := cas.StoreFile(ctx, projID, userID, "spec_v2_dup.pdf", "standard", testData, "application/pdf")
	if err != nil {
		t.Fatalf("store file 2 to db failed: %v", err)
	}
	if !ref2.IsDeduplicated {
		t.Errorf("expected second upload to be deduplicated in DB")
	}

	// 3. 存储并检索向量切片（内容带唯一后缀，避免跨次运行 chunk_hash 冲突）
	chunk, err := cas.StoreMemoryChunk(ctx, MemoryChunkInput{
		Domain:           "project_doc",
		ProjectID:        projID,
		UserID:           userID,
		Title:            "光机指标要求",
		Content:          "内窥镜视场角不低于 100 度，光学畸变应严格控制在 5% 以内。 run=" + time.Now().Format(time.RFC3339Nano),
		ImportanceWeight: 1.2,
	})
	if err != nil {
		t.Fatalf("store memory chunk to db failed: %v", err)
	}
	if chunk.ID <= 0 {
		t.Errorf("expected positive chunk ID from db")
	}

	// 4. 检索向量切片
	hits, err := cas.SearchMemory(ctx, MemorySearchOptions{
		Domain:            "project_doc",
		ProjectID:         projID,
		QueryText:         "视场角",
		ApplyTimeDecay:    true,
		DecayHalfLifeDays: 7.0,
		TopK:              5,
	})
	if err != nil {
		t.Fatalf("search memory in db failed: %v", err)
	}
	if len(hits) == 0 {
		t.Fatalf("expected memory hits from db, got 0")
	}
	if hits[0].Chunk.Title != "光机指标要求" {
		t.Errorf("unexpected top hit title from db: %s", hits[0].Chunk.Title)
	}
}