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;
///
/// 为 RAG 向量集成测试创建并销毁独立 PostgreSQL + pgvector 数据库。
/// 遵循 PostgresWorkflowFixture 模式:随机命名数据库 → 建扩展 + 建表 → 测试 → 清理。
/// pgvector 扩展必须在 EnsureCreatedAsync 之前创建(UseVector 依赖该扩展)。
/// embedding vector(768) 列、HNSW 索引和 Content 列 trigram GIN 索引在 EnsureCreated 后通过 raw SQL 补充
/// (与 AddPgTrgmIndex 迁移保持一致——fixture 走 EnsureCreated 不执行迁移)。
///
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!;
/// 当前测试数据库连接串。
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();
}
}
/// 创建指向独立测试数据库的 RagDbContext(带 AuditInterceptor)。
public RagDbContext CreateDbContext(IInterceptor[]? interceptors = null)
{
var optionsBuilder = new DbContextOptionsBuilder()
.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();
}
}