diff --git a/core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs b/core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs index c6532c8500..b5bba60d7a 100644 --- a/core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs +++ b/core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs @@ -6,6 +6,7 @@ using System.Text; using Azure.Mcp.Tests.Client.Attributes; using Azure.Mcp.Tests.Client.Helpers; +using Azure.Mcp.Tests.Generated; using Azure.Mcp.Tests.Generated.Models; using Azure.Mcp.Tests.Helpers; using Xunit; @@ -18,6 +19,8 @@ public abstract class RecordedCommandTestsBase(ITestOutputHelper output, TestPro protected TestProxy? Proxy { get; private set; } = fixture.Proxy; + protected virtual RecordingOptions? RecordingOptions { get; private set; } = null; + protected string RecordingId { get; private set; } = string.Empty; /// @@ -152,6 +155,18 @@ public override async ValueTask InitializeAsync() // apply custom matcher if test has attribute await ApplyAttributeMatcherSettings(); + + SetRecordingOptions(RecordingOptions); + } + + public void SetRecordingOptions(RecordingOptions? options) + { + if (Proxy == null || TestMode != TestMode.Live || options == null) + { + return; + } + + Proxy.AdminClient.SetRecordingOptions(options, RecordingId); } private async Task ApplyAttributeMatcherSettings() @@ -221,10 +236,7 @@ public async Task StartProxyAsync(TestProxyFixture fixture) Proxy = fixture.Proxy; // onetime on starting the proxy, we have initialized the livetest settings so lets add some additional sanitizers by default - if (EnableDefaultSanitizerAdditions) - { - PopulateDefaultSanitizers(); - } + PopulateDefaultSanitizers(); // onetime registration of default sanitizers // and deregistering default sanitizers that we don't want diff --git a/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs b/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs index 7c109a0614..2772df9661 100644 --- a/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs +++ b/tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs @@ -12,10 +12,11 @@ namespace Azure.Mcp.Tools.Cosmos.Services; -public class CosmosService(ISubscriptionService subscriptionService, ITenantService tenantService, ICacheService cacheService) +public class CosmosService(ISubscriptionService subscriptionService, ITenantService tenantService, IHttpClientFactory httpClientFactory, ICacheService cacheService) : BaseAzureService(tenantService), ICosmosService, IDisposable { private readonly ISubscriptionService _subscriptionService = subscriptionService ?? throw new ArgumentNullException(nameof(subscriptionService)); + private readonly IHttpClientFactory _httpClientFactory = httpClientFactory ?? throw new ArgumentNullException(nameof(httpClientFactory)); private readonly ICacheService _cacheService = cacheService ?? throw new ArgumentNullException(nameof(cacheService)); private const string CosmosBaseUri = "https://{0}.documents.azure.com:443/"; private const string CacheGroup = "cosmos"; @@ -65,6 +66,8 @@ private async Task CreateCosmosClientWithAuth( clientOptions.MaxRetryWaitTimeOnRateLimitedRequests = TimeSpan.FromSeconds(retryPolicy.MaxDelaySeconds); } + clientOptions.HttpClientFactory = () => _httpClientFactory.CreateClient(); + CosmosClient cosmosClient; switch (authMethod) { diff --git a/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosCommandTests.cs b/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosCommandTests.cs index 914606d75d..354a9cec3c 100644 --- a/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosCommandTests.cs +++ b/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosCommandTests.cs @@ -4,14 +4,28 @@ using System.Text.Json; using Azure.Mcp.Tests; using Azure.Mcp.Tests.Client; +using Azure.Mcp.Tests.Client.Helpers; +using Azure.Mcp.Tests.Generated.Models; using Xunit; namespace Azure.Mcp.Tools.Cosmos.LiveTests; -public class CosmosCommandTests(ITestOutputHelper output) - : CommandTestsBase(output), - IClassFixture +public class CosmosCommandTests(ITestOutputHelper output, TestProxyFixture fixture) : RecordedCommandTestsBase(output, fixture) { + protected override RecordingOptions? RecordingOptions => new() + { + HandleRedirects = false + }; + + /// + /// 3493 = $..name + /// + public override List DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "3493"]; + + public override CustomDefaultMatcher? TestMatcher => new() + { + IgnoredHeaders = "x-ms-activity-id,x-ms-cosmos-correlated-activityid" + }; [Fact] public async Task Should_list_storage_accounts_by_subscription_id() @@ -66,12 +80,13 @@ public async Task Should_list_cosmos_database_containers_by_database_name() [Fact] public async Task Should_query_cosmos_database_container_items() { + var resourceBaseName = TestMode == Tests.Helpers.TestMode.Playback ? "Sanitized" : Settings.ResourceBaseName; var result = await CallToolAsync( "cosmos_database_container_item_query", new() { { "subscription", Settings.SubscriptionId }, - { "account", Settings.ResourceBaseName }, + { "account", resourceBaseName }, { "database", "ToDoList" }, { "container", "Items" } }); @@ -99,12 +114,13 @@ public async Task Should_list_cosmos_accounts() [Fact] public async Task Should_show_single_item_from_cosmos_account() { + var resourceBaseName = TestMode == Tests.Helpers.TestMode.Playback ? "Sanitized" : Settings.ResourceBaseName; var dbResult = await CallToolAsync( "cosmos_database_list", new() { { "subscription", Settings.SubscriptionId }, - { "account", Settings.ResourceBaseName } + { "account", resourceBaseName } } ); var databases = dbResult.AssertProperty("databases"); @@ -114,11 +130,7 @@ public async Task Should_show_single_item_from_cosmos_account() // The agent will choose one, for this test we're going to take the first one var firstDatabase = dbEnum.First(); - string dbName = firstDatabase.ValueKind == JsonValueKind.String - ? firstDatabase.GetString()! - : firstDatabase.ValueKind == JsonValueKind.Object - ? firstDatabase.GetProperty("name").GetString()! - : throw new InvalidOperationException($"Unexpected database element ValueKind: {firstDatabase.ValueKind}"); + string dbName = RegisterOrRetrieveVariable("database", GetStringOrNameElementString(firstDatabase, "database")); Assert.False(string.IsNullOrEmpty(dbName)); var containerResult = await CallToolAsync( @@ -126,8 +138,8 @@ public async Task Should_show_single_item_from_cosmos_account() new() { { "subscription", Settings.SubscriptionId }, - { "account", Settings.ResourceBaseName }, - { "database", dbName! } + { "account", resourceBaseName }, + { "database", dbName } }); var containers = containerResult.AssertProperty("containers"); Assert.Equal(JsonValueKind.Array, containers.ValueKind); @@ -136,11 +148,7 @@ public async Task Should_show_single_item_from_cosmos_account() // The agent will choose one, for this test we're going to take the first one var firstContainer = contEnum.First(); - string containerName = firstContainer.ValueKind == JsonValueKind.String - ? firstContainer.GetString()! - : firstContainer.ValueKind == JsonValueKind.Object - ? firstContainer.GetProperty("name").GetString()! - : throw new InvalidOperationException($"Unexpected container element ValueKind: {firstContainer.ValueKind}"); + string containerName = RegisterOrRetrieveVariable("container", GetStringOrNameElementString(firstContainer, "container")); Assert.False(string.IsNullOrEmpty(containerName)); var itemResult = await CallToolAsync( @@ -148,8 +156,8 @@ public async Task Should_show_single_item_from_cosmos_account() new() { { "subscription", Settings.SubscriptionId }, - { "account", Settings.ResourceBaseName }, - { "database", dbName! }, + { "account", resourceBaseName }, + { "database", dbName }, { "container", containerName! } }); var items = itemResult.AssertProperty("items"); @@ -160,48 +168,69 @@ public async Task Should_show_single_item_from_cosmos_account() [Fact] public async Task Should_list_and_query_multiple_databases_and_containers() { + var resourceBaseName = TestMode == Tests.Helpers.TestMode.Playback ? "Sanitized" : Settings.ResourceBaseName; var dbResult = await CallToolAsync( "cosmos_database_list", new() { { "subscription", Settings.SubscriptionId }, - { "account", Settings.ResourceBaseName } + { "account", resourceBaseName } } ); var databases = dbResult.AssertProperty("databases"); Assert.Equal(JsonValueKind.Array, databases.ValueKind); var databasesEnum = databases.EnumerateArray(); Assert.True(databasesEnum.Any()); + var dbNumber = 0; foreach (var db in databasesEnum) { - string dbName = db.ValueKind == JsonValueKind.String - ? db.GetString()! - : db.ValueKind == JsonValueKind.Object - ? db.GetProperty("name").GetString()! - : throw new InvalidOperationException($"Unexpected database element ValueKind: {db.ValueKind}"); + string dbName = RegisterOrRetrieveVariable("database" + dbNumber, GetStringOrNameElementString(db, "database")); Assert.False(string.IsNullOrEmpty(dbName)); var containerResult = await CallToolAsync( "cosmos_database_container_list", - new() { { "subscription", Settings.SubscriptionId }, { "account", Settings.ResourceBaseName! }, { "database", dbName! } }); + new() { + { "subscription", Settings.SubscriptionId }, + { "account", resourceBaseName }, + { "database", dbName } + }); var containers = containerResult.AssertProperty("containers"); Assert.Equal(JsonValueKind.Array, containers.ValueKind); var contEnum = containers.EnumerateArray(); + var containerNumber = 0; foreach (var container in contEnum) { - string containerName = container.ValueKind == JsonValueKind.String - ? container.GetString()! - : container.ValueKind == JsonValueKind.Object - ? container.GetProperty("name").GetString()! - : throw new InvalidOperationException($"Unexpected container element ValueKind: {container.ValueKind}"); + string containerName = RegisterOrRetrieveVariable("db" + dbNumber + "/container" + containerNumber, GetStringOrNameElementString(container, "container")); Assert.False(string.IsNullOrEmpty(containerName)); var itemResult = await CallToolAsync( "cosmos_database_container_item_query", - new() { { "subscription", Settings.SubscriptionId }, { "account", Settings.ResourceBaseName! }, { "database", dbName! }, { "container", containerName! } }); + new() { + { "subscription", Settings.SubscriptionId }, + { "account", resourceBaseName }, + { "database", dbName }, + { "container", containerName } + }); var items = itemResult.AssertProperty("items"); Assert.Equal(JsonValueKind.Array, items.ValueKind); + containerNumber++; } + dbNumber++; } } + + private static string GetStringOrNameElementString(JsonElement element, string propertyName) + { + if (element.ValueKind == JsonValueKind.String) + { + return element.GetString()!; + } + + if (element.ValueKind == JsonValueKind.Object) + { + return element.GetProperty("name").GetString()!; + } + + throw new InvalidOperationException($"Unexpected {propertyName} element ValueKind: {element.ValueKind}"); + } } diff --git a/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosDbFixture.cs b/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosDbFixture.cs deleted file mode 100644 index ac4f28084d..0000000000 --- a/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosDbFixture.cs +++ /dev/null @@ -1,61 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -using System.Text.Json; -using System.Text.Json.Serialization; -using Azure.Mcp.Core.Services.Azure.Authentication; -using Azure.Mcp.Tests.Client.Helpers; -using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.Logging.Abstractions; -using Xunit; - -namespace Azure.Mcp.Tools.Cosmos.LiveTests; - -public class CosmosDbFixture : IAsyncLifetime -{ - private CosmosClient? _client; - - public async ValueTask InitializeAsync() - { - var settingsFixture = new LiveTestSettingsFixture(); - await settingsFixture.InitializeAsync(); - - var tokenProvider = new SingleIdentityTokenCredentialProvider(NullLoggerFactory.Instance); - - _client = new CosmosClient( - accountEndpoint: $"https://{settingsFixture.Settings.ResourceBaseName}.documents.azure.com:443/", - tokenCredential: await tokenProvider.GetTokenCredentialAsync(default, default) - ); - Container container = _client.GetContainer("ToDoList", "Items"); - ToDoItem entry = new ToDoItem(); - Stream stream = ToStream(entry); - await container.UpsertItemStreamAsync(streamPayload: stream, partitionKey: new PartitionKey(entry.id) - ); - } - - public ValueTask DisposeAsync() - { - _client?.Dispose(); - return ValueTask.CompletedTask; - } - - private static Stream ToStream(ToDoItem input) - { - MemoryStream stream = new(); - JsonSerializer.Serialize(stream, input, JsonContext.Default.ToDoItem); - stream.Seek(0, SeekOrigin.Begin); - return stream; - } -} - -[JsonSourceGenerationOptions(WriteIndented = true)] -[JsonSerializable(typeof(ToDoItem))] -[JsonSerializable(typeof(List))] -public partial class JsonContext : JsonSerializerContext { } - -public class ToDoItem -{ - public string id { get; set; } = Guid.NewGuid().ToString(); - public string title { get; set; } = "Test Task"; - public bool completed { get; set; } = false; -} diff --git a/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/assets.json b/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/assets.json new file mode 100644 index 0000000000..929b71e01b --- /dev/null +++ b/tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/assets.json @@ -0,0 +1,6 @@ +{ + "AssetsRepo": "Azure/azure-sdk-assets", + "AssetsRepoPrefixPath": "", + "TagPrefix": "Azure.Mcp.Tools.Cosmos.LiveTests", + "Tag": "Azure.Mcp.Tools.Cosmos.LiveTests_3821fa67e8" +}