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

using FluentAssertions;
using PiiRedaction.Core.Detection;
using PiiRedaction.Core.Models;
using PiiRedaction.Infrastructure.Onnx;

namespace PiiRedaction.Infrastructure.Tests.Onnx;

[TestFixture]
public sealed class RoutingOnnxNerModelRunnerTests
{
    [Test]
    public void PredictEntities_LatinOnly_UsesEnglishRunnerOnly()
    {
        var english = new FakeLanguageNerRunner("Ravi Kumar");
        var tamil = new FakeLanguageNerRunner("தமிழ் பெயர்");
        var router = new RoutingOnnxNerModelRunner(english, tamil, enableTamilNer: true);

        var entities = router.PredictEntities("Customer Ravi Kumar called.");

        entities.Should().ContainSingle(entity => entity.Value == "Ravi Kumar");
        english.CallCount.Should().Be(1);
        tamil.CallCount.Should().Be(0);
    }

    [Test]
    public void PredictEntities_TamilOnly_UsesTamilRunnerOnly()
    {
        var english = new FakeLanguageNerRunner("Ravi Kumar");
        var tamil = new FakeLanguageNerRunner("ராஜேஷ்");
        var router = new RoutingOnnxNerModelRunner(english, tamil, enableTamilNer: true);

        var entities = router.PredictEntities("வாடிக்கையாளர் ராஜேஷ்");

        entities.Should().ContainSingle(entity => entity.Value == "ராஜேஷ்");
        english.CallCount.Should().Be(0);
        tamil.CallCount.Should().Be(1);
    }

    [Test]
    public void PredictEntities_Mixed_InvokesBothRunners()
    {
        var english = new FakeLanguageNerRunner("EnglishName");
        var tamil = new FakeLanguageNerRunner("தமிழ்");
        var router = new RoutingOnnxNerModelRunner(english, tamil, enableTamilNer: true);

        router.PredictEntities("Rajesh மற்றும் Priya");

        english.CallCount.Should().Be(1);
        tamil.CallCount.Should().Be(1);
    }

    [Test]
    public void PredictEntities_NoLetters_InvokesNeither()
    {
        var english = new FakeLanguageNerRunner("ignored");
        var tamil = new FakeLanguageNerRunner("ignored");
        var router = new RoutingOnnxNerModelRunner(english, tamil, enableTamilNer: true);

        router.PredictEntities("9876543210").Should().BeEmpty();
        english.CallCount.Should().Be(0);
        tamil.CallCount.Should().Be(0);
    }

    [Test]
    public void PredictEntities_TamilDisabled_SkipsTamilRunnerForMixedText()
    {
        var english = new FakeLanguageNerRunner("EnglishName");
        var tamil = new FakeLanguageNerRunner("தமிழ்");
        var router = new RoutingOnnxNerModelRunner(english, tamil, enableTamilNer: false);

        router.PredictEntities("Rajesh மற்றும் Priya");

        english.CallCount.Should().Be(1);
        tamil.CallCount.Should().Be(0);
    }

    [Test]
    public void MergePersonSpans_PrefersLongerOverlappingSpan()
    {
        var entities = new[]
        {
            CreatePerson("Raj", 0, 3),
            CreatePerson("Rajesh", 0, 6)
        };

        var merged = RoutingOnnxNerModelRunner.MergePersonSpans(entities);

        merged.Should().ContainSingle(entity => entity.Value == "Rajesh");
    }

    private static PiiEntity CreatePerson(string value, int start, int length) =>
        new(PiiEntityType.Person, value, start, length, PiiDetectionSource.Ner);

    private sealed class FakeLanguageNerRunner : IOnnxNerModelRunner
    {
        private readonly string _personValue;

        public FakeLanguageNerRunner(string personValue) => _personValue = personValue;

        public int CallCount { get; private set; }

        public bool IsModelAvailable => true;

        public IReadOnlyList<PiiEntity> PredictEntities(string text)
        {
            CallCount++;
            return
            [
                new PiiEntity(
                    PiiEntityType.Person,
                    _personValue,
                    0,
                    _personValue.Length,
                    PiiDetectionSource.Ner)
            ];
        }
    }
}