llm-pii-poc/tests/PiiRedaction.Core.Tests/Detection/OnnxNerPiiDetectorTests.cs

using FluentAssertions;
using PiiRedaction.Core.Detection;
using PiiRedaction.Core.Models;
using PiiRedaction.Core.Tests.TestSupport;

namespace PiiRedaction.Core.Tests.Detection;

[TestFixture]
public sealed class OnnxNerPiiDetectorTests
{
    [Test]
    public void Detect_ModelAvailable_ReturnsRunnerEntities()
    {
        var runner = new FakeOnnxNerModelRunner
        {
            IsModelAvailable = true,
            EntitiesToReturn =
            [
                new PiiEntity(PiiEntityType.Person, "Onnx Person", 0, 11, PiiDetectionSource.Ner)
            ]
        };

        var detector = new OnnxNerPiiDetector(runner);
        var entities = detector.Detect("Any text");

        entities.Should().ContainSingle(entity => entity.Value == "Onnx Person");
        runner.LastPredictedText.Should().Be("Any text");
    }

    [Test]
    public void Detect_ModelUnavailable_ReturnsEmpty()
    {
        var detector = CreateDetector(modelAvailable: false);
        var entities = detector.Detect("Customer Ravi Kumar with email test@x.com.");

        entities.Should().BeEmpty();
    }

    [TestCase(null)]
    [TestCase("")]
    [TestCase("   ")]
    public void Detect_InvalidInput_ThrowsArgumentException(string? text)
    {
        var detector = CreateDetector(modelAvailable: false);
        var action = () => detector.Detect(text!);
        action.Should().Throw<ArgumentException>();
    }

    [Test]
    public void Detect_ModelAvailable_RecordsInvokedModels()
    {
        var runner = new FakeOnnxNerModelRunner
        {
            IsModelAvailable = true,
            InvokedModelsToReturn = [NerModelOrigin.English, NerModelOrigin.Tamil]
        };

        var detector = new OnnxNerPiiDetector(runner);
        detector.Detect("Mixed prompt");

        detector.LastInvokedModels.Should().Equal(NerModelOrigin.English, NerModelOrigin.Tamil);
    }

    [Test]
    public void Detect_ModelUnavailable_ClearsInvokedModels()
    {
        var runner = new FakeOnnxNerModelRunner
        {
            IsModelAvailable = false,
            InvokedModelsToReturn = [NerModelOrigin.English]
        };

        var detector = new OnnxNerPiiDetector(runner);
        detector.Detect("Any text");

        detector.LastInvokedModels.Should().BeEmpty();
    }

    private static OnnxNerPiiDetector CreateDetector(bool modelAvailable)
    {
        var runner = new FakeOnnxNerModelRunner { IsModelAvailable = modelAvailable };
        return new OnnxNerPiiDetector(runner);
    }
}