124 lines
5.1 KiB
C#
124 lines
5.1 KiB
C#
using Microsoft.EntityFrameworkCore;
|
||
using Microsoft.EntityFrameworkCore.Diagnostics;
|
||
using Npgsql;
|
||
using Pgvector.EntityFrameworkCore;
|
||
using BuildFlow.Infrastructure.Persistence;
|
||
using BuildFlow.Shared.Infrastructure;
|
||
using Xunit;
|
||
|
||
namespace BuildFlow.Rag.Tests;
|
||
|
||
/// <summary>
|
||
/// 为 RAG 向量集成测试创建并销毁独立 PostgreSQL + pgvector 数据库。
|
||
/// 遵循 PostgresWorkflowFixture 模式：随机命名数据库 → 建扩展 + 建表 → 测试 → 清理。
|
||
/// pgvector 扩展必须在 EnsureCreatedAsync 之前创建（UseVector 依赖该扩展）。
|
||
/// embedding vector(768) 列、HNSW 索引和 Content 列 trigram GIN 索引在 EnsureCreated 后通过 raw SQL 补充
|
||
/// （与 AddPgTrgmIndex 迁移保持一致——fixture 走 EnsureCreated 不执行迁移）。
|
||
/// </summary>
|
||
public sealed class PostgresRagFixture : IAsyncLifetime
|
||
{
|
||
private const string DefaultConnection =
|
||
"Host=127.0.0.1;Port=55432;Username=buildflow;Password=buildflow;Database=postgres";
|
||
|
||
private string _databaseName = default!;
|
||
private string _adminConnectionString = default!;
|
||
|
||
/// <summary>当前测试数据库连接串。</summary>
|
||
public string ConnectionString { get; private set; } = default!;
|
||
|
||
public async Task InitializeAsync()
|
||
{
|
||
var configured = Environment.GetEnvironmentVariable("BUILDFLOW_RAG_TEST_POSTGRES_CONNECTION")
|
||
?? DefaultConnection;
|
||
|
||
var adminBuilder = new NpgsqlConnectionStringBuilder(configured)
|
||
{
|
||
Database = "postgres",
|
||
Pooling = false,
|
||
};
|
||
_adminConnectionString = adminBuilder.ConnectionString;
|
||
_databaseName = $"rag_vector_{Guid.NewGuid():N}";
|
||
|
||
// 创建独立数据库
|
||
await using (var connection = new NpgsqlConnection(_adminConnectionString))
|
||
{
|
||
await connection.OpenAsync();
|
||
await using var command = connection.CreateCommand();
|
||
command.CommandText = $"CREATE DATABASE \"{_databaseName}\"";
|
||
await command.ExecuteNonQueryAsync();
|
||
}
|
||
|
||
var databaseBuilder = new NpgsqlConnectionStringBuilder(configured)
|
||
{
|
||
Database = _databaseName,
|
||
Pooling = true,
|
||
};
|
||
ConnectionString = databaseBuilder.ConnectionString;
|
||
|
||
// pgvector 扩展必须在 EnsureCreated 之前安装
|
||
await using (var connection = new NpgsqlConnection(ConnectionString))
|
||
{
|
||
await connection.OpenAsync();
|
||
await using var cmd = connection.CreateCommand();
|
||
cmd.CommandText = """
|
||
CREATE EXTENSION IF NOT EXISTS vector;
|
||
CREATE EXTENSION IF NOT EXISTS pg_trgm;
|
||
""";
|
||
await cmd.ExecuteNonQueryAsync();
|
||
}
|
||
|
||
// EnsureCreated 建出 b_ 前缀的表结构
|
||
await using var db = CreateDbContext();
|
||
await db.Database.EnsureCreatedAsync();
|
||
|
||
// 补充 embedding 列和 HNSW 索引（EF 不映射这些）
|
||
await using (var connection = new NpgsqlConnection(ConnectionString))
|
||
{
|
||
await connection.OpenAsync();
|
||
await using var cmd = connection.CreateCommand();
|
||
cmd.CommandText = """
|
||
ALTER TABLE b_document_chunks ADD COLUMN IF NOT EXISTS embedding vector(768);
|
||
CREATE INDEX IF NOT EXISTS ix_b_document_chunks_embedding_hnsw
|
||
ON b_document_chunks USING hnsw (embedding vector_cosine_ops);
|
||
CREATE INDEX IF NOT EXISTS ix_b_document_chunks_content_trgm
|
||
ON b_document_chunks USING gin ("Content" gin_trgm_ops);
|
||
""";
|
||
await cmd.ExecuteNonQueryAsync();
|
||
}
|
||
}
|
||
|
||
/// <summary>创建指向独立测试数据库的 RagDbContext（带 AuditInterceptor）。</summary>
|
||
public RagDbContext CreateDbContext(IInterceptor[]? interceptors = null)
|
||
{
|
||
var optionsBuilder = new DbContextOptionsBuilder<RagDbContext>()
|
||
.UseNpgsql(ConnectionString, o => o.UseVector())
|
||
.ConfigureWarnings(w => w.Ignore(RelationalEventId.PendingModelChangesWarning));
|
||
|
||
if (interceptors is { Length: > 0 })
|
||
optionsBuilder.AddInterceptors(interceptors);
|
||
else
|
||
optionsBuilder.AddInterceptors(new AuditInterceptor(new Stub.TestCurrentUserContext()));
|
||
|
||
return new RagDbContext(optionsBuilder.Options);
|
||
}
|
||
|
||
public async Task DisposeAsync()
|
||
{
|
||
NpgsqlConnection.ClearAllPools();
|
||
await using var connection = new NpgsqlConnection(_adminConnectionString);
|
||
await connection.OpenAsync();
|
||
|
||
await using (var terminate = connection.CreateCommand())
|
||
{
|
||
terminate.CommandText =
|
||
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = @database AND pid <> pg_backend_pid()";
|
||
terminate.Parameters.AddWithValue("database", _databaseName);
|
||
await terminate.ExecuteNonQueryAsync();
|
||
}
|
||
|
||
await using var drop = connection.CreateCommand();
|
||
drop.CommandText = $"DROP DATABASE IF EXISTS \"{_databaseName}\"";
|
||
await drop.ExecuteNonQueryAsync();
|
||
}
|
||
}
|