llm-pii-poc/tests/PiiRedaction.Core.Tests/Integration/RealNerPipelineTests.cs

using FluentAssertions;

using PiiRedaction.Core.Abstractions;

using PiiRedaction.Core.Models;

using PiiRedaction.Core.Tests.TestSupport;

using PiiRedaction.Tests.Shared;



namespace PiiRedaction.Core.Tests.Integration;



/// <summary>

/// End-to-end pipeline proof using the real ONNX NER model (no fakes).

/// </summary>

[TestFixture]

[Category("RealModel")]

public sealed class RealNerPipelineTests : RealNerModelFixture

{

    private IPromptSanitizer _sanitizer = null!;



    [OneTimeSetUp]

    public void OneTimeSetUpPipeline()

    {

        _sanitizer = ProductionPipelineFactory.CreateWithRealModel(Runner);

    }



    [Test]

    public void Sanitize_FullFinancialWithCustomer_RedactsAllPiiTypes()

    {

        const string prompt =

            "Customer Ravi Kumar with email ravi.kumar@gmail.com and phone 9876543210 has LoanNumber LN-456789 and PAN ABCDE1234F. Please summarize this customer issue.";



        var result = _sanitizer.Sanitize(new SanitizationRequest(prompt));



        result.SanitizedPrompt.Should().Be(

            "Customer <PERSON_1> with email <EMAIL_1> and phone <PHONE_1> has LoanNumber <LOAN_NUMBER_1> and PAN <PAN_1>. Please summarize this customer issue.");



        result.SanitizedPrompt.Should().NotContainAny("Ravi Kumar", "ravi.kumar@gmail.com", "9876543210", "LN-456789", "ABCDE1234F");

        result.DetectedEntities.Should().Contain(entity =>

            entity.Type == PiiEntityType.Person && entity.Source == PiiDetectionSource.Ner);

    }



    [Test]

    public void Sanitize_TwoCustomersInOnePrompt_RedactsBothPeople()

    {

        const string prompt = "Customer Ravi Kumar and Customer Priya Nair disputed the same charge.";



        var result = _sanitizer.Sanitize(new SanitizationRequest(prompt));



        result.SanitizedPrompt.Should().Contain("<PERSON_1>");

        result.SanitizedPrompt.Should().Contain("<PERSON_2>");

        result.SanitizedPrompt.Should().NotContainAny("Ravi Kumar", "Priya Nair");

        result.DetectedEntities.Count(entity => entity.Type == PiiEntityType.Person).Should().BeGreaterThanOrEqualTo(2);

    }



    [Test]

    public void Sanitize_NoPiiCleanTicket_PassesThroughUnchanged()

    {

        const string prompt = "What is the status of ticket TKT-99887 and when will the API maintenance end?";



        var result = _sanitizer.Sanitize(new SanitizationRequest(prompt));



        result.SanitizedPrompt.Should().Be(prompt);

        result.DetectedEntities.Should().BeEmpty();

    }

}