meeting-assistant/MeetingAssistant.Tests/FunAsrStreamingTranscriptionProviderTests.cs

using System.Text.Json;
using MeetingAssistant.Recording;
using MeetingAssistant.Transcription;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Tests;

public sealed class FunAsrStreamingTranscriptionProviderTests
{
    [Fact]
    public async Task ProviderStreamsPcmChunksAndEmitsSpeakerAttributedSegments()
    {
        var connection = new FakeFunAsrWebSocketConnection(
            """
            {"mode":"2pass-offline","wav_name":"meeting","text":"hello from mike","is_final":true,"spk_name":"Mike","spk_score":0.92,"timestamp":[[1000,1800]]}
            """);
        var provider = new FunAsrStreamingTranscriptionProvider(
            new FakeFunAsrWebSocketConnectionFactory(connection),
            new NoopFunAsrBackendLifecycle(),
            Options.Create(new MeetingAssistantOptions
            {
                FunAsr = new FunAsrOptions
                {
                    Endpoint = "ws://localhost:10095",
                    Mode = "2pass",
                    ChunkSize = [5, 10, 5],
                    ChunkInterval = 10,
                    FinalResultTimeout = TimeSpan.FromSeconds(1)
                }
            }),
            NullLogger<FunAsrStreamingTranscriptionProvider>.Instance);
        var chunks = new[]
        {
            new AudioChunk([1, 0, 2, 0], 16000, 1),
            new AudioChunk([3, 0, 4, 0], 16000, 1)
        };

        var segments = await CollectAsync(provider.TranscribeAsync(ToAsyncEnumerable(chunks), SpeechRecognitionPipelineOptions.Default, CancellationToken.None));

        Assert.Equal(new Uri("ws://localhost:10095"), connection.ConnectedEndpoint);
        Assert.Equal(2, connection.BinaryMessages.Count);
        Assert.Equal(chunks[0].Pcm, connection.BinaryMessages[0]);
        Assert.Equal(chunks[1].Pcm, connection.BinaryMessages[1]);
        Assert.Equal("Mike", Assert.Single(segments).Speaker);
        Assert.Equal("hello from mike", segments[0].Text);
        Assert.Equal(TimeSpan.FromSeconds(1), segments[0].Start);
        Assert.Equal(TimeSpan.FromMilliseconds(1800), segments[0].End);

        using var firstMessage = JsonDocument.Parse(connection.TextMessages[0]);
        Assert.Equal("2pass", firstMessage.RootElement.GetProperty("mode").GetString());
        Assert.Equal("pcm", firstMessage.RootElement.GetProperty("wav_format").GetString());
        Assert.Equal(16000, firstMessage.RootElement.GetProperty("audio_fs").GetInt32());
        Assert.True(firstMessage.RootElement.GetProperty("is_speaking").GetBoolean());

        using var finalMessage = JsonDocument.Parse(connection.TextMessages[^1]);
        Assert.False(finalMessage.RootElement.GetProperty("is_speaking").GetBoolean());
    }

    [Fact]
    public async Task ProviderUsesSentenceSpeakerFieldsAndUnknownFallback()
    {
        var connection = new FakeFunAsrWebSocketConnection(
            """
            {"mode":"2pass-offline","wav_name":"meeting","text":"ignored","is_final":true,"sentence_info":[{"text":"first sentence","spk":0,"start":100,"end":600},{"text":"second sentence","start":700,"end":1200}]}
            """);
        var provider = new FunAsrStreamingTranscriptionProvider(
            new FakeFunAsrWebSocketConnectionFactory(connection),
            new NoopFunAsrBackendLifecycle(),
            Options.Create(new MeetingAssistantOptions
            {
                FunAsr = new FunAsrOptions { FinalResultTimeout = TimeSpan.FromSeconds(1) }
            }),
            NullLogger<FunAsrStreamingTranscriptionProvider>.Instance);

        var segments = await CollectAsync(provider.TranscribeAsync(
            ToAsyncEnumerable([new AudioChunk([1, 0], 16000, 1)]),
            SpeechRecognitionPipelineOptions.Default,
            CancellationToken.None));

        Assert.Collection(
            segments,
            first =>
            {
                Assert.Equal("Speaker 0", first.Speaker);
                Assert.Equal("first sentence", first.Text);
                Assert.Equal(TimeSpan.FromMilliseconds(100), first.Start);
                Assert.Equal(TimeSpan.FromMilliseconds(600), first.End);
            },
            second =>
            {
                Assert.Equal("Unknown", second.Speaker);
                Assert.Equal("second sentence", second.Text);
                Assert.Equal(TimeSpan.FromMilliseconds(700), second.Start);
                Assert.Equal(TimeSpan.FromMilliseconds(1200), second.End);
            });
    }

    [Fact]
    public void BuildHotwordsSerializesDictationWordsWithDeduplication()
    {
        var hotwords = FunAsrStreamingTranscriptionProvider.BuildHotwords([
            "PBI",
            "pbi",
            "Product Backlog Item",
            " "
        ]);

        using var document = JsonDocument.Parse(hotwords);
        Assert.Equal(20, document.RootElement.GetProperty("PBI").GetInt32());
        Assert.Equal(20, document.RootElement.GetProperty("Product Backlog Item").GetInt32());
        Assert.Equal(2, document.RootElement.EnumerateObject().Count());
    }

    private static async IAsyncEnumerable<AudioChunk> ToAsyncEnumerable(IEnumerable<AudioChunk> chunks)
    {
        foreach (var chunk in chunks)
        {
            yield return chunk;
            await Task.Yield();
        }
    }

    private static async Task<List<TranscriptionSegment>> CollectAsync(IAsyncEnumerable<TranscriptionSegment> segments)
    {
        var collected = new List<TranscriptionSegment>();
        await foreach (var segment in segments)
        {
            collected.Add(segment);
        }

        return collected;
    }

    private sealed class FakeFunAsrWebSocketConnectionFactory : IFunAsrWebSocketConnectionFactory
    {
        private readonly FakeFunAsrWebSocketConnection connection;

        public FakeFunAsrWebSocketConnectionFactory(FakeFunAsrWebSocketConnection connection)
        {
            this.connection = connection;
        }

        public Task<IFunAsrWebSocketConnection> ConnectAsync(Uri endpoint, CancellationToken cancellationToken)
        {
            connection.ConnectedEndpoint = endpoint;
            return Task.FromResult<IFunAsrWebSocketConnection>(connection);
        }
    }

    private sealed class NoopFunAsrBackendLifecycle : IFunAsrBackendLifecycle
    {
        public Task EnsureStartedAsync(CancellationToken cancellationToken)
        {
            return Task.CompletedTask;
        }
    }

    private sealed class FakeFunAsrWebSocketConnection : IFunAsrWebSocketConnection
    {
        private readonly Queue<string> responses;

        public FakeFunAsrWebSocketConnection(params string[] responses)
        {
            this.responses = new Queue<string>(responses);
        }

        public Uri? ConnectedEndpoint { get; set; }

        public List<string> TextMessages { get; } = [];

        public List<byte[]> BinaryMessages { get; } = [];

        public Task SendTextAsync(string message, CancellationToken cancellationToken)
        {
            TextMessages.Add(message);
            return Task.CompletedTask;
        }

        public Task SendBinaryAsync(ReadOnlyMemory<byte> message, CancellationToken cancellationToken)
        {
            BinaryMessages.Add(message.ToArray());
            return Task.CompletedTask;
        }

        public Task<string?> ReceiveTextAsync(CancellationToken cancellationToken)
        {
            return Task.FromResult(responses.Count == 0 ? null : responses.Dequeue());
        }

        public Task CloseOutputAsync(CancellationToken cancellationToken)
        {
            return Task.CompletedTask;
        }

        public ValueTask DisposeAsync()
        {
            return ValueTask.CompletedTask;
        }
    }
}