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

using System.Security.Cryptography;
using BuildFlow.Infrastructure.Obsidian;
using BuildFlow.Shared.Domain.Errors;
using BuildFlow.Shared.Domain.Exceptions;
using FluentAssertions;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using Xunit;

namespace BuildFlow.Rag.Tests;

/// <summary>
/// Obsidian 服务端目录访问和容量边界测试。
/// </summary>
[Collection(ObsidianVaultScannerCollection.Name)]
public sealed class ObsidianVaultScannerTests : IDisposable
{
    private readonly string _testRoot = Path.Combine(
        Path.GetTempPath(),
        $"buildflow-obsidian-{Guid.NewGuid():N}");
    private readonly string _allowedRoot;
    private readonly string _outsideRoot;

    public ObsidianVaultScannerTests()
    {
        _allowedRoot = Directory.CreateDirectory(Path.Combine(_testRoot, "allowed")).FullName;
        _outsideRoot = Directory.CreateDirectory(Path.Combine(_testRoot, "outside")).FullName;
    }

    [Fact]
    public async Task OutsideAllowedRootShouldBeRejected()
    {
        var scanner = CreateScanner(_allowedRoot);

        var act = () => CollectAsync(scanner.ScanAsync(_outsideRoot));

        await act.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.InvalidParam);
    }

    [Fact]
    public async Task ParentDirectoryTraversalShouldBeRejected()
    {
        var traversal = Path.Combine(_allowedRoot, "child", "..", "..", "outside");
        var scanner = CreateScanner(_allowedRoot);

        var act = () => CollectAsync(scanner.ScanAsync(traversal));

        await act.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.InvalidParam);
    }

    [Fact]
    public async Task SymlinkEscapingAllowedRootShouldBeRejected()
    {
        var link = Path.Combine(_allowedRoot, "escape");
        Directory.CreateSymbolicLink(link, _outsideRoot);
        await File.WriteAllTextAsync(Path.Combine(_outsideRoot, "secret.md"), "secret");
        var scanner = CreateScanner(_allowedRoot);

        var act = () => CollectAsync(scanner.ScanAsync(_allowedRoot));

        await act.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.InvalidParam);
    }

    [Fact]
    public async Task FileCountLimitShouldUseTooLargeCode()
    {
        await File.WriteAllTextAsync(Path.Combine(_allowedRoot, "1.md"), "1");
        await File.WriteAllTextAsync(Path.Combine(_allowedRoot, "2.md"), "2");
        var scanner = CreateScanner(_allowedRoot, maxFiles: 1);

        var act = () => CollectAsync(scanner.ScanAsync(_allowedRoot));

        await act.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.TooLarge);
    }

    [Fact]
    public async Task SingleAndTotalSizeLimitsShouldUseTooLargeCode()
    {
        await File.WriteAllBytesAsync(Path.Combine(_allowedRoot, "large.md"), new byte[32]);
        var singleScanner = CreateScanner(_allowedRoot, maxFileBytes: 16);
        var singleAct = () => CollectAsync(singleScanner.ScanAsync(_allowedRoot));
        await singleAct.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.TooLarge);

        await File.WriteAllBytesAsync(Path.Combine(_allowedRoot, "small.md"), new byte[8]);
        var totalScanner = CreateScanner(
            _allowedRoot,
            maxFileBytes: 64,
            maxTotalBytes: 32);
        var totalAct = () => CollectAsync(totalScanner.ScanAsync(_allowedRoot));
        await totalAct.Should().ThrowAsync<BusinessException>()
            .Where(x => x.Code == ErrorCode.TooLarge);
    }

    [Fact]
    public async Task ValidVaultShouldReturnRelativePathSizeAndStreamingHash()
    {
        var nested = Directory.CreateDirectory(Path.Combine(_allowedRoot, "Daily"));
        var path = Path.Combine(nested.FullName, "note.md");
        await File.WriteAllTextAsync(path, "hello");

        var files = await CollectAsync(CreateScanner(_allowedRoot).ScanAsync(_allowedRoot));

        files.Should().ContainSingle();
        files[0].RelativePath.Should().Be("Daily/note.md");
        files[0].FileSize.Should().Be(5);
        files[0].ContentHash.Should().Be(
            Convert.ToHexString(SHA256.HashData("hello"u8)).ToLowerInvariant());
    }

    [Fact]
    public async Task LargeVaultShouldNotAllocateTheCombinedFileSize()
    {
        for (var i = 0; i < 8; i++)
        {
            await using var stream = File.Create(Path.Combine(_allowedRoot, $"{i}.md"));
            stream.SetLength(8L * 1024 * 1024);
        }

        GC.Collect();
        GC.WaitForPendingFinalizers();
        GC.Collect();
        var baseline = GC.GetTotalMemory(forceFullCollection: true);
        var peak = baseline;
        var count = 0;

        await foreach (var _ in CreateScanner(
                           _allowedRoot,
                           maxFileBytes: 9L * 1024 * 1024,
                           maxTotalBytes: 70L * 1024 * 1024).ScanAsync(_allowedRoot))
        {
            count++;
            peak = Math.Max(peak, GC.GetTotalMemory(forceFullCollection: false));
        }

        count.Should().Be(8);
        (peak - baseline).Should().BeLessThan(32L * 1024 * 1024);
    }

    private static ObsidianVaultScanner CreateScanner(
        string allowedRoot,
        int maxFiles = 100,
        long maxFileBytes = 16L * 1024 * 1024,
        long maxTotalBytes = 128L * 1024 * 1024)
    {
        var options = Options.Create(new ObsidianOptions
        {
            AllowedRoots = [allowedRoot],
            MaxFiles = maxFiles,
            MaxFileBytes = maxFileBytes,
            MaxTotalBytes = maxTotalBytes,
        });
        return new ObsidianVaultScanner(
            options,
            NullLogger<ObsidianVaultScanner>.Instance);
    }

    private static async Task<List<ObsidianVaultFile>> CollectAsync(
        IAsyncEnumerable<ObsidianVaultFile> source)
    {
        var result = new List<ObsidianVaultFile>();
        await foreach (var item in source) result.Add(item);
        return result;
    }

    public void Dispose()
    {
        if (Directory.Exists(_testRoot)) Directory.Delete(_testRoot, recursive: true);
    }
}