build-flow.api/tests/BuildFlow.Rag.Tests/PostgresRagFixture.cs

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 索引在 EnsureCreated 后通过 raw SQL 补充。
/// </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";
            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);
                """;
            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();
    }
}