Skip to content

Migrate Cosmos to recordings #1402

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alzimmermsft wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Migrate Cosmos to recordings #1402

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c792762
Migrate Cosmos to recordings
alzimmermsft Dec 17, 2025
4e56e0e
Merge branch 'main' into Recordings_Cosmos
scbedd Dec 17, 2025
65f6e7e
commit recordings with redirect fix
scbedd Dec 18, 2025
5fc2268
Merge branch 'main' into Recordings_Cosmos
scbedd Dec 19, 2025
e142b81
repair tests manually. filing issue for follow-up
scbedd Dec 19, 2025
9cdf6aa
format!
scbedd Dec 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

/// <summary>
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion tools/Azure.Mcp.Tools.Cosmos/src/Services/CosmosService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -65,6 +66,8 @@ private async Task<CosmosClient> CreateCosmosClientWithAuth(
clientOptions.MaxRetryWaitTimeOnRateLimitedRequests = TimeSpan.FromSeconds(retryPolicy.MaxDelaySeconds);
}

clientOptions.HttpClientFactory = () => _httpClientFactory.CreateClient();

CosmosClient cosmosClient;
switch (authMethod)
{
Expand Down
93 changes: 61 additions & 32 deletions tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosCommandTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<CosmosDbFixture>
public class CosmosCommandTests(ITestOutputHelper output, TestProxyFixture fixture) : RecordedCommandTestsBase(output, fixture)
{
protected override RecordingOptions? RecordingOptions => new()
{
HandleRedirects = false
};

/// <summary>
/// 3493 = $..name
/// </summary>
public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "3493"];
Comment on lines +21 to +23
Copy link

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sanitizer ID should be "AZSDK3493" instead of "3493". All other tools in the codebase consistently use the "AZSDK" prefix for this sanitizer (e.g., AcrCommandTests, AppConfigCommandTests, QuotaCommandTests, SearchCommandTests, SqlCommandTests, SignalRCommandTests).

Suggested change
/// 3493 = $..name
/// </summary>
public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "3493"];
/// AZSDK3493 = $..name
/// </summary>
public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "AZSDK3493"];

Copilot uses AI. Check for mistakes.
Comment on lines +21 to +23
Copy link

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment should reference "AZSDK3493" instead of "3493" to match the correct sanitizer ID. This is consistent with how other tools document this sanitizer (e.g., "AZSDK3493 = $..name" in AcrCommandTests, AppConfigCommandTests, QuotaCommandTests, etc.).

Suggested change
/// 3493 = $..name
/// </summary>
public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "3493"];
/// AZSDK3493 = $..name
/// </summary>
public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "AZSDK3493"];

Copilot uses AI. Check for mistakes.

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()
Expand Down Expand Up @@ -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" }
});
Expand Down Expand Up @@ -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");
Expand All @@ -114,20 +130,16 @@ 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(
"cosmos_database_container_list",
new()
{
{ "subscription", Settings.SubscriptionId },
{ "account", Settings.ResourceBaseName },
{ "database", dbName! }
{ "account", resourceBaseName },
{ "database", dbName }
});
var containers = containerResult.AssertProperty("containers");
Assert.Equal(JsonValueKind.Array, containers.ValueKind);
Expand All @@ -136,20 +148,16 @@ 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(
"cosmos_database_container_item_query",
new()
{
{ "subscription", Settings.SubscriptionId },
{ "account", Settings.ResourceBaseName },
{ "database", dbName! },
{ "account", resourceBaseName },
{ "database", dbName },
{ "container", containerName! }
Copy link

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The null-forgiving operator (!) is unnecessary here. The variable 'containerName' is declared as a non-nullable string, and there's an assertion on line 152 that verifies it's not null or empty. The null-forgiving operator should be removed for cleaner code.

Suggested change
{ "container", containerName! }
{ "container", containerName }

Copilot uses AI. Check for mistakes.
});
var items = itemResult.AssertProperty("items");
Expand All @@ -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}");
}
}
61 changes: 0 additions & 61 deletions tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/CosmosDbFixture.cs
View file
Open in desktop

This file was deleted.

6 changes: 6 additions & 0 deletions tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/assets.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"AssetsRepo": "Azure/azure-sdk-assets",
"AssetsRepoPrefixPath": "",
"TagPrefix": "Azure.Mcp.Tools.Cosmos.LiveTests",
"Tag": "Azure.Mcp.Tools.Cosmos.LiveTests_3821fa67e8"
}