chore: import upstream snapshot with attribution

dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/A2AServerServiceCollectionExtensionsTests.cs

@@ -0,0 +1,459 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using A2A;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using Moq;
+using Moq.Protected;
+
+namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests;
+
+/// <summary>
+/// Unit tests for the <see cref="A2AServerServiceCollectionExtensions"/> class.
+/// </summary>
+public sealed class A2AServerServiceCollectionExtensionsTests
+{
+    /// <summary>
+    /// Verifies that AddA2AServer with an agent name registers a keyed A2AServer
+    /// that can be resolved from the service provider.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithAgentName_ResolvesKeyedA2AServerAsync()
+    {
+        // Arrange
+        const string AgentName = "test-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        // Act
+        services.AddA2AServer(AgentName);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer with an agent instance registers a keyed A2AServer
+    /// that can be resolved from the service provider using the agent's name.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithAgentInstance_ResolvesKeyedA2AServerAsync()
+    {
+        // Arrange
+        const string AgentName = "instance-agent";
+        var agentMock = CreateAgentMock(AgentName);
+        var services = new ServiceCollection();
+
+        // Act
+        services.AddA2AServer(agentMock.Object);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that when no ITaskStore or AgentSessionStore are registered,
+    /// AddA2AServer falls back to noop session store default and resolves successfully.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithNoCustomStores_FallsBackToNoopSessionStoreDefaultAsync()
+    {
+        // Arrange
+        const string AgentName = "default-stores-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        // Act
+        services.AddA2AServer(AgentName);
+
+        // Assert - resolution succeeds without any stores registered
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that when a custom ITaskStore is registered, AddA2AServer uses it
+    /// instead of the default InMemoryTaskStore.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithCustomTaskStore_ResolvesSuccessfullyAsync()
+    {
+        // Arrange
+        const string AgentName = "custom-taskstore-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        var mockTaskStore = new Mock<ITaskStore>();
+        services.AddKeyedSingleton(AgentName, mockTaskStore.Object);
+
+        // Act
+        services.AddA2AServer(AgentName);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that when a custom AgentSessionStore is registered, AddA2AServer uses it
+    /// instead of the default InMemoryAgentSessionStore.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithCustomAgentSessionStore_ResolvesSuccessfullyAsync()
+    {
+        // Arrange
+        const string AgentName = "custom-sessionstore-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        var mockSessionStore = new Mock<AgentSessionStore>();
+        services.AddKeyedSingleton(AgentName, mockSessionStore.Object);
+
+        // Act
+        services.AddA2AServer(AgentName);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that when a custom IAgentHandler is registered, AddA2AServer uses it
+    /// instead of creating a default A2AAgentHandler.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithCustomAgentHandler_ResolvesSuccessfullyAsync()
+    {
+        // Arrange
+        const string AgentName = "custom-handler-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        var mockHandler = new Mock<IAgentHandler>();
+        services.AddKeyedSingleton(AgentName, mockHandler.Object);
+
+        // Act
+        services.AddA2AServer(AgentName);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that the configureOptions callback is invoked when provided.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithConfigureOptions_InvokesCallbackAsync()
+    {
+        // Arrange
+        const string AgentName = "options-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        bool callbackInvoked = false;
+
+        // Act
+        services.AddA2AServer(AgentName, options =>
+        {
+            callbackInvoked = true;
+            options.AgentRunMode = AgentRunMode.AllowBackgroundIfSupported;
+        });
+
+        // Assert - callback is invoked during resolution
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+        Assert.True(callbackInvoked);
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer with a null configureOptions does not throw.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithNullConfigureOptions_ResolvesSuccessfullyAsync()
+    {
+        // Arrange
+        const string AgentName = "null-options-agent";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        // Act
+        services.AddA2AServer(AgentName, configureOptions: null);
+
+        // Assert
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetKeyedService<A2AServer>(AgentName);
+        Assert.NotNull(server);
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer throws when the agent name is null.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithNullAgentName_ThrowsArgumentException()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+
+        // Act & Assert
+        Assert.ThrowsAny<ArgumentException>(() => services.AddA2AServer(agentName: null!));
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer throws when the agent name is whitespace.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithWhitespaceAgentName_ThrowsArgumentException()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+
+        // Act & Assert
+        Assert.ThrowsAny<ArgumentException>(() => services.AddA2AServer(agentName: "  "));
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer throws when the services parameter is null.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithNullServices_ThrowsArgumentNullException()
+    {
+        // Arrange
+        IServiceCollection services = null!;
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => services.AddA2AServer("agent"));
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer with an agent instance throws when the agent is null.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithNullAgent_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => services.AddA2AServer(agent: null!));
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer with an agent instance throws when the agent's Name is null.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithAgent_NullName_ThrowsArgumentNullException()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        var agentMock = new Mock<AIAgent>();
+        agentMock.Setup(a => a.Name).Returns((string?)null);
+
+        // Act & Assert
+        ArgumentNullException exception = Assert.Throws<ArgumentNullException>(() =>
+            services.AddA2AServer(agentMock.Object));
+
+        Assert.Equal("agent.Name", exception.ParamName);
+    }
+
+    /// <summary>
+    /// Verifies that AddA2AServer with an agent instance throws when the agent's Name is whitespace.
+    /// </summary>
+    [Fact]
+    public void AddA2AServer_WithAgent_WhitespaceName_ThrowsArgumentException()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        var agentMock = new Mock<AIAgent>();
+        agentMock.Setup(a => a.Name).Returns("   ");
+
+        // Act & Assert
+        ArgumentException exception = Assert.Throws<ArgumentException>(() =>
+            services.AddA2AServer(agentMock.Object));
+
+        Assert.Equal("agent.Name", exception.ParamName);
+    }
+
+    /// <summary>
+    /// Verifies that when a custom <see cref="IAgentHandler"/> is registered as a keyed service,
+    /// the <see cref="A2AServer"/> uses it to process requests instead of the default handler.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithCustomHandler_CustomHandlerIsInvokedOnRequestAsync()
+    {
+        // Arrange
+        const string AgentName = "custom-handler-wiring";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        var mockHandler = new Mock<IAgentHandler>();
+        mockHandler
+            .Setup(h => h.ExecuteAsync(
+                It.IsAny<RequestContext>(),
+                It.IsAny<AgentEventQueue>(),
+                It.IsAny<CancellationToken>()))
+            .Returns((RequestContext _, AgentEventQueue eq, CancellationToken ct) =>
+                eq.EnqueueMessageAsync(
+                    new Message { MessageId = "resp", Role = Role.Agent, Parts = [new Part { Text = "Reply" }] }, ct).AsTask());
+
+        services.AddKeyedSingleton(AgentName, mockHandler.Object);
+
+        services.AddA2AServer(AgentName);
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetRequiredKeyedService<A2AServer>(AgentName);
+
+        // Act
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+        var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token);
+
+        // Assert - the custom handler was invoked, not the default A2AAgentHandler
+        mockHandler.Verify(
+            h => h.ExecuteAsync(
+                It.IsAny<RequestContext>(),
+                It.IsAny<AgentEventQueue>(),
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+        Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase);
+        Assert.NotNull(response.Message);
+    }
+
+    /// <summary>
+    /// Verifies that when a custom <see cref="AgentSessionStore"/> is registered as a keyed service
+    /// and no custom <see cref="IAgentHandler"/> is registered, the default handler uses the custom
+    /// session store for session management during request processing.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithCustomSessionStore_NoHandler_SessionStoreIsUsedOnRequestAsync()
+    {
+        // Arrange
+        const string AgentName = "custom-sessionstore-wiring";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMock(AgentName).Object);
+
+        var mockSessionStore = new Mock<AgentSessionStore>();
+        mockSessionStore
+            .Setup(x => x.GetSessionAsync(
+                It.IsAny<AIAgent>(),
+                It.IsAny<string>(),
+                It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new TestAgentSession());
+        mockSessionStore
+            .Setup(x => x.SaveSessionAsync(
+                It.IsAny<AIAgent>(),
+                It.IsAny<string>(),
+                It.IsAny<AgentSession>(),
+                It.IsAny<CancellationToken>()))
+            .Returns(ValueTask.CompletedTask);
+
+        services.AddKeyedSingleton(AgentName, mockSessionStore.Object);
+
+        services.AddA2AServer(AgentName);
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetRequiredKeyedService<A2AServer>(AgentName);
+
+        // Act
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+        var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token);
+
+        // Assert - the custom session store was used, not InMemoryAgentSessionStore
+        mockSessionStore.Verify(
+            x => x.GetSessionAsync(
+                It.IsAny<AIAgent>(),
+                It.IsAny<string>(),
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+        Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase);
+        Assert.NotNull(response.Message);
+    }
+
+    /// <summary>
+    /// Verifies that when no custom stores or handlers are registered, the server uses
+    /// the default noop session store and processes requests successfully end-to-end.
+    /// </summary>
+    [Fact]
+    public async Task AddA2AServer_WithNoCustomStores_DefaultStoresProcessRequestSuccessfullyAsync()
+    {
+        // Arrange
+        const string AgentName = "default-stores-request";
+        var services = new ServiceCollection();
+        services.AddKeyedSingleton(AgentName, (_, _) => CreateAgentMockForRequests(AgentName).Object);
+
+        services.AddA2AServer(AgentName);
+        await using var provider = services.BuildServiceProvider();
+        var server = provider.GetRequiredKeyedService<A2AServer>(AgentName);
+
+        // Act
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
+        var response = await server.SendMessageAsync(CreateTestSendMessageRequest(), cts.Token);
+
+        // Assert - request was processed successfully with default noop session store
+        Assert.NotNull(response);
+        Assert.Equal(SendMessageResponseCase.Message, response.PayloadCase);
+        Assert.NotNull(response.Message);
+    }
+
+    private static SendMessageRequest CreateTestSendMessageRequest() =>
+        new()
+        {
+            Message = new Message
+            {
+                MessageId = "test-id",
+                Role = Role.User,
+                Parts = [new Part { Text = "Hello" }]
+            }
+        };
+
+    private static Mock<AIAgent> CreateAgentMock(string name)
+    {
+        Mock<AIAgent> agentMock = new() { CallBase = true };
+        agentMock.SetupGet(x => x.Name).Returns(name);
+        agentMock
+            .Protected()
+            .Setup<ValueTask<AgentSession>>("CreateSessionCoreAsync", ItExpr.IsAny<CancellationToken>())
+            .ReturnsAsync(new TestAgentSession());
+        agentMock
+            .Protected()
+            .Setup<Task<AgentResponse>>("RunCoreAsync",
+                ItExpr.IsAny<IEnumerable<ChatMessage>>(),
+                ItExpr.IsAny<AgentSession?>(),
+                ItExpr.IsAny<AgentRunOptions?>(),
+                ItExpr.IsAny<CancellationToken>())
+            .ReturnsAsync(new AgentResponse([new ChatMessage(ChatRole.Assistant, "Test response")]));
+
+        return agentMock;
+    }
+
+    /// <summary>
+    /// Creates a mock <see cref="AIAgent"/> with session serialization support, suitable for
+    /// tests that exercise the full request processing path with <see cref="InMemoryAgentSessionStore"/>.
+    /// </summary>
+    private static Mock<AIAgent> CreateAgentMockForRequests(string name)
+    {
+        Mock<AIAgent> agentMock = CreateAgentMock(name);
+        agentMock
+            .Protected()
+            .Setup<ValueTask<JsonElement>>("SerializeSessionCoreAsync",
+                ItExpr.IsAny<AgentSession>(),
+                ItExpr.IsAny<JsonSerializerOptions?>(),
+                ItExpr.IsAny<CancellationToken>())
+            .ReturnsAsync(JsonDocument.Parse("{}").RootElement);
+
+        return agentMock;
+    }
+
+    private sealed class TestAgentSession : AgentSession;
+}