File: Model\AIAssistant\AssistantChatDataContextTests.cs
Web Access
Project: src\tests\Aspire.Dashboard.Tests\Aspire.Dashboard.Tests.csproj (Aspire.Dashboard.Tests)
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Threading.Channels;
using Aspire.Dashboard.Model;
using Aspire.Dashboard.Model.Assistant;
using Aspire.Dashboard.Otlp.Model;
using Aspire.Dashboard.Otlp.Storage;
using Aspire.Dashboard.Tests.Integration.Playwright.Infrastructure;
using Aspire.Dashboard.Tests.Shared;
using Aspire.Tests.Shared.DashboardModel;
using Google.Protobuf.Collections;
using OpenTelemetry.Proto.Logs.V1;
using OpenTelemetry.Proto.Trace.V1;
using Xunit;
using static Aspire.Tests.Shared.Telemetry.TelemetryTestHelpers;
 
namespace Aspire.Dashboard.Tests.Model.AIAssistant;
 
public class AssistantChatDataContextTests
{
    private static readonly DateTime s_testTime = new(1970, 1, 1, 0, 0, 0, DateTimeKind.Utc);
 
    [Fact]
    public void GetLimitFromEndWithSummary_UnderLimits_ReturnAll()
    {
        // Arrange
        var values = new List<string>();
        for (var i = 0; i < 10; i++)
        {
            values.Add(new string((char)('a' + i), 16));
        }
 
        // Act
        var (items, message) = AssistantChatDataContext.GetLimitFromEndWithSummary(values, totalValues: values.Count, limit: 20, "test item", s => s, s => ((string)s).Length);
 
        // Assert
        Assert.Equal(10, items.Count);
        Assert.Equal("Returned 10 test items.", message);
    }
 
    [Fact]
    public void GetLimitFromEndWithSummary_UnderTotal_ReturnPassedIn()
    {
        // Arrange
        var values = new List<string>();
        for (var i = 0; i < 10; i++)
        {
            values.Add(new string((char)('a' + i), 16));
        }
 
        // Act
        var (items, message) = AssistantChatDataContext.GetLimitFromEndWithSummary(values, totalValues: 100, limit: 20, "test item", s => s, s => ((string)s).Length);
 
        // Assert
        Assert.Equal(10, items.Count);
        Assert.Equal("Returned latest 10 test items. Earlier 90 test items not returned because of size limits.", message);
    }
 
    [Fact]
    public void GetLimitFromEndWithSummary_ExceedCountLimit_ReturnMostRecentItems()
    {
        // Arrange
        var values = new List<string>();
        for (var i = 0; i < 10; i++)
        {
            values.Add(new string((char)('a' + i), 2));
        }
 
        // Act
        var (items, message) = AssistantChatDataContext.GetLimitFromEndWithSummary(values, totalValues: 100, limit: 5, "test item", s => s, s => ((string)s).Length);
 
        // Assert
        Assert.Collection(items,
            s => Assert.Equal("ff", s),
            s => Assert.Equal("gg", s),
            s => Assert.Equal("hh", s),
            s => Assert.Equal("ii", s),
            s => Assert.Equal("jj", s));
        Assert.Equal("Returned latest 5 test items. Earlier 95 test items not returned because of size limits.", message);
    }
 
    [Fact]
    public void GetLimitFromEndWithSummary_ExceedTokenLimit_ReturnMostRecentItems()
    {
        const int textLength = 1024 * 2;
 
        // Arrange
        var values = new List<string>();
        for (var i = 0; i < 10; i++)
        {
            values.Add(new string((char)('a' + i), textLength));
        }
 
        // Act
        var (items, message) = AssistantChatDataContext.GetLimitFromEndWithSummary(values, limit: 10, "test item", s => s, s => ((string)s).Length);
 
        // Assert
        Assert.Collection(items,
            s => Assert.Equal(new string('g', textLength), s),
            s => Assert.Equal(new string('h', textLength), s),
            s => Assert.Equal(new string('i', textLength), s),
            s => Assert.Equal(new string('j', textLength), s));
        Assert.Equal("Returned latest 4 test items. Earlier 6 test items not returned because of size limits.", message);
    }
 
    [Fact]
    public async Task GetStructuredLogs_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var repository = CreateRepository();
 
        var scopeLogs = new ScopeLogs();
        for (var i = 0; i < 20; i++)
        {
            var logRecord = CreateLogRecord(message: $"Log {i}: {new string((char)('a' + i), 10_000)}", time: s_testTime.AddMinutes(i));
            scopeLogs.LogRecords.Add(logRecord);
        }
        var addContext = new AddContext();
        repository.AddLogs(addContext, new RepeatedField<ResourceLogs>()
        {
            new ResourceLogs
            {
                Resource = CreateResource(),
                ScopeLogs = { scopeLogs }
            }
        });
        var dataContext = CreateAssistantChatDataContext(telemetryRepository: repository);
 
        // Act
        var result = await dataContext.GetStructuredLogsAsync(resourceName: null, CancellationToken.None);
 
        // Assert
        for (var i = 6; i < 20; i++)
        {
            Assert.Contains($"Log {i}:", result);
        }
        Assert.Contains("Returned latest 14 log entries. Earlier 6 log entries not returned because of size limits.", result);
    }
 
    [Fact]
    public async Task GetTraces_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var repository = CreateRepository();
 
        var scopeSpans = new ScopeSpans();
        for (var i = 0; i < 20; i++)
        {
            var span = CreateSpan(traceId: $"{i}", spanId: $"{i}-1", startTime: s_testTime.AddMinutes(i), endTime: s_testTime.AddMinutes(10), attributes: [new KeyValuePair<string, string>("message", $"Log {i}: {new string((char)('a' + i), 10_000)}")]);
            scopeSpans.Spans.Add(span);
        }
        var addContext = new AddContext();
        repository.AddTraces(addContext, new RepeatedField<ResourceSpans>()
        {
            new ResourceSpans
            {
                Resource = CreateResource(),
                ScopeSpans = { scopeSpans }
            }
        });
        var dataContext = CreateAssistantChatDataContext(telemetryRepository: repository);
 
        // Act
        var result = await dataContext.GetTracesAsync(resourceName: null, CancellationToken.None);
 
        // Assert
        for (var i = 7; i < 20; i++)
        {
            Assert.Contains($"Test span. Id: {i}", result);
        }
        Assert.Contains("Returned latest 13 traces. Earlier 7 traces not returned because of size limits.", result);
    }
 
    [Fact]
    public async Task GetConsoleLogs_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var consoleLogsChannel = Channel.CreateUnbounded<IReadOnlyList<ResourceLogLine>>();
        var testResource = ModelTestHelpers.CreateResource(resourceName: "test-resource", state: KnownResourceState.Running);
        var dashboardClient = new TestDashboardClient(
            isEnabled: true,
            consoleLogsChannelProvider: name =>
            {
                return consoleLogsChannel;
            },
            initialResources: [testResource]);
 
        var dataContext = CreateAssistantChatDataContext(dashboardClient: dashboardClient);
 
        for (var i = 0; i < 20; i++)
        {
            var line = new string((char)('a' + i), 10_000);
            consoleLogsChannel.Writer.TryWrite([new ResourceLogLine(i + 1, line, IsErrorMessage: false)]);
        }
        consoleLogsChannel.Writer.Complete();
 
        // Act
        var result = await dataContext.GetConsoleLogsAsync(resourceName: "test-resource", CancellationToken.None);
 
        // Assert
        for (var i = 5; i < 20; i++)
        {
            var line = AIHelpers.LimitLength(new string((char)('a' + i), 10_000));
            Assert.Contains(line, result);
        }
        Assert.Contains("Returned latest 15 console logs. Earlier 5 console logs not returned because of size limits.", result);
    }
 
    internal static AssistantChatDataContext CreateAssistantChatDataContext(TelemetryRepository? telemetryRepository = null, IDashboardClient? dashboardClient = null)
    {
        var context = new AssistantChatDataContext(
            telemetryRepository ?? CreateRepository(),
            dashboardClient ?? new MockDashboardClient(),
            [],
            new TestStringLocalizer<Dashboard.Resources.AIAssistant>());
 
        return context;
    }
}