llm-pii-poc/tests/PiiRedaction.Infrastructure.Tests/Onnx/RealNerModelRunnerTests.cs

using FluentAssertions;

using PiiRedaction.Core.Models;

using PiiRedaction.Tests.Shared;



namespace PiiRedaction.Infrastructure.Tests.Onnx;



[TestFixture]

[Category("RealModel")]

public sealed class RealNerModelRunnerTests : RealNerModelFixture

{

    [TestCase("Customer Ravi Kumar called about billing.", "Ravi", "Ravi Kumar")]

    [TestCase("Mr. John Smith called about a duplicate debit.", "John", "John Smith")]

    public void PredictEntities_DetectsPersonWithCorrectSpan(string prompt, string expectedNamePart, string expectedValue)

    {

        var result = Runner.PredictEntities(prompt);
        var entities = result.Entities;

        entities.Should().Contain(entity =>
            entity.Type == PiiEntityType.Person &&
            entity.Source == PiiDetectionSource.Ner &&
            entity.ModelOrigin == NerModelOrigin.English &&
            entity.Value.Contains(expectedNamePart, StringComparison.Ordinal) &&

            prompt.AsSpan(entity.StartIndex, entity.Length).ToString() == entity.Value);



        entities.Should().Contain(entity => entity.Value == expectedValue);
        result.InvokedModels.Should().Equal(NerModelOrigin.English);
    }
}