llm-pii-poc/tests/PiiRedaction.Core.Tests/TestSupport/FakeOnnxNerModelRunner.cs

using PiiRedaction.Core.Detection;
using PiiRedaction.Core.Models;

namespace PiiRedaction.Core.Tests.TestSupport;

public sealed class FakeOnnxNerModelRunner : IOnnxNerModelRunner
{
    public bool IsModelAvailable { get; set; }

    public IReadOnlyList<PiiEntity> EntitiesToReturn { get; set; } = [];

    public IReadOnlyList<NerModelOrigin> InvokedModelsToReturn { get; set; } = [NerModelOrigin.English];

    public string? LastPredictedText { get; private set; }

    public NerPredictionResult PredictEntities(string text)
    {
        LastPredictedText = text;
        return new NerPredictionResult(EntitiesToReturn, InvokedModelsToReturn);
    }
}