.NET: chore: support retries on Cosmos storage creation (#402)

* support retries

* tests + registration options

* fix ordering ..

* HK + update packages

* fix paths

* Update dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosTestFixture.cs

* re create project and fix some pk usage

* fix all tests

* try workflow?

* wip 1

* fix definition

* try with cosmos_use_emulator env?

* try ignore SSL errors?

* other cert verifications

* hardcode to 8081?

* proper valuation of ENV

* logging

* ensure db exsists for CI

* bump

* cleanup

* fix usage

* nit comment

* try only release for stability?

* try skip some flaky tests

* merge fixes + rollback container

* reimplement with iasyncdisposable pattern

* remove example doc struct
This commit is contained in:
Korolev Dmitry
2025-08-21 20:37:12 +02:00
committed by GitHub
Unverified
parent 1116916737
commit 25291de8cb
22 changed files with 436 additions and 110 deletions
@@ -0,0 +1,158 @@
#
# This workflow runs Cosmos DB integration tests using the Cosmos DB emulator.
#
name: dotnet-cosmosdb-integration-tests
on:
workflow_dispatch:
pull_request:
branches: ["main", "feature*"]
paths:
- dotnet/tests/CosmosDB.IntegrationTests/**
- dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/**
- '.github/workflows/dotnet-cosmosdb-integration-tests.yml'
merge_group:
branches: ["main"]
push:
branches: ["main", "feature*"]
paths:
- dotnet/tests/CosmosDB.IntegrationTests/**
- dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/**
- '.github/workflows/dotnet-cosmosdb-integration-tests.yml'
schedule:
- cron: "0 2 * * *" # Run at 2 AM UTC daily
env:
COSMOSDB_TESTS_USE_EMULATOR_CICD: "true"
jobs:
build-and-test:
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
include:
- { targetFramework: "net9.0", os: "ubuntu-latest", configuration: Release }
# - { targetFramework: "net9.0", os: "ubuntu-latest", configuration: Debug }
services:
cosmosdb:
image: mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest
ports:
- 8081:8081
env:
AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE: "false"
AZURE_COSMOS_EMULATOR_PARTITION_COUNT: "20" # the more the better for stable tests
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
sparse-checkout: |
.
.github
dotnet
- name: Setup dotnet
uses: actions/setup-dotnet@v4.3.1
with:
global-json-file: ${{ github.workspace }}/dotnet/global.json
- name: Build dotnet solutions
shell: bash
run: |
export SOLUTIONS=$(find ./dotnet/ -type f -name "*.slnx" | tr '\n' ' ')
for solution in $SOLUTIONS; do
dotnet build $solution -c ${{ matrix.configuration }} --warnaserror
done
- name: Package install check
shell: bash
# All frameworks are only built for the release configuration, so we only run this step for the release configuration
# and dotnet new doesn't support net472
if: matrix.configuration == 'Release' && matrix.targetFramework != 'net472'
run: |
TEMP_DIR=$(mktemp -d)
export SOLUTIONS=$(find ./dotnet/ -type f -name "*.slnx" | tr '\n' ' ')
for solution in $SOLUTIONS; do
dotnet pack $solution /property:TargetFrameworks=${{ matrix.targetFramework }} -c ${{ matrix.configuration }} --no-build --no-restore --output "$TEMP_DIR/artifacts"
done
pushd "$TEMP_DIR"
# Create a new console app to test the package installation
dotnet new console -f ${{ matrix.targetFramework }} --name packcheck --output consoleapp
# Create minimal nuget.config and use only dotnet nuget commands
echo '<?xml version="1.0" encoding="utf-8"?><configuration><packageSources><clear /></packageSources></configuration>' > consoleapp/nuget.config
# Add sources with local first using dotnet nuget commands
dotnet nuget add source ../artifacts --name local --configfile consoleapp/nuget.config
dotnet nuget add source https://api.nuget.org/v3/index.json --name nuget.org --configfile consoleapp/nuget.config
# Change to project directory to ensure local nuget.config is used
pushd consoleapp
dotnet add packcheck.csproj package Microsoft.Extensions.AI.Agents --prerelease
dotnet build -f ${{ matrix.targetFramework }} -c ${{ matrix.configuration }} packcheck.csproj
# Clean up
popd
popd
rm -rf "$TEMP_DIR"
- name: Wait for Cosmos DB Emulator to be ready
run: |
set -e
for i in $(seq 1 120); do
if curl -sk https://localhost:8081/_explorer/emulator.pem -o /dev/null; then
echo "Emulator is up."
break
fi
echo "Waiting for emulator... ($i/120)"
sleep 2
done
- name: Install emulator TLS certificate into system trust store
run: |
set -e
sudo apt-get update
sudo apt-get install -y ca-certificates curl openssl
# Fetch the PEM directly from the emulator's explorer endpoint
curl -sk https://localhost:8081/_explorer/emulator.pem -o cosmos-emulator.crt
# Install with the correct .crt extension so update-ca-certificates picks it up
sudo cp cosmos-emulator.crt /usr/local/share/ca-certificates/cosmos-emulator.crt
sudo update-ca-certificates
- name: Verify TLS now trusts the emulator
run: |
# Use -servername to avoid SNI warning and check verification
echo | openssl s_client -connect localhost:8081 -servername localhost 2>/dev/null | grep -E "Verify return code|subject=|issuer="
# Expect: "Verify return code: 0 (ok)"
- name: Run Cosmos DB Integration Tests
shell: bash
run: |
# Run the specific CosmosDB integration tests
dotnet test ./dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj \
-f ${{ matrix.targetFramework }} \
-c ${{ matrix.configuration }} \
--no-build \
-v Normal \
--logger trx \
--collect:"XPlat Code Coverage" \
--results-directory:"TestResults/Coverage/" \
-- DataCollectionRunSettings.DataCollectors.DataCollector.Configuration.ExcludeByAttribute=GeneratedCodeAttribute,CompilerGeneratedAttribute,ExcludeFromCodeCoverageAttribute
# Generate test reports and check coverage
- name: Generate test reports
uses: danielpalme/ReportGenerator-GitHub-Action@5.4.11
with:
reports: "./TestResults/Coverage/**/coverage.cobertura.xml"
targetdir: "./TestResults/Reports"
reporttypes: "HtmlInline;JsonSummary"
- name: Upload coverage report artifact
uses: actions/upload-artifact@v4
with:
name: CosmosDB-CoverageReport-${{ matrix.os }}-${{ matrix.targetFramework }}-${{ matrix.configuration }}
path: ./TestResults/Reports
+11 -6
View File
@@ -5,14 +5,18 @@
<ManagePackageVersionsCentrally>true</ManagePackageVersionsCentrally>
<CentralPackageTransitivePinningEnabled>true</CentralPackageTransitivePinningEnabled>
</PropertyGroup>
<PropertyGroup>
<!-- Aspire -->
<AspireAppHostSdkVersion>9.4.1</AspireAppHostSdkVersion>
</PropertyGroup>
<ItemGroup>
<!-- Azure.* -->
<PackageVersion Include="Aspire.Azure.AI.OpenAI" Version="9.3.1-preview.1.25305.6" />
<PackageVersion Include="Aspire.Hosting.AppHost" Version="9.4.0" />
<PackageVersion Include="Aspire.Hosting.Azure.CognitiveServices" Version="9.4.0" />
<PackageVersion Include="Aspire.Hosting.Azure.CosmosDB" Version="9.4.0" />
<PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="9.3.1" />
<PackageVersion Include="Aspire.Hosting.Testing" Version="9.3.1" />
<PackageVersion Include="Aspire.Azure.AI.OpenAI" Version="9.4.1-preview.1.25408.4" />
<PackageVersion Include="Aspire.Hosting.AppHost" Version="9.4.1" />
<PackageVersion Include="Aspire.Hosting.Azure.CognitiveServices" Version="9.4.1" />
<PackageVersion Include="Aspire.Hosting.Azure.CosmosDB" Version="9.4.1" />
<PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="9.4.1" />
<PackageVersion Include="Aspire.Hosting.Testing" Version="9.4.1" />
<PackageVersion Include="Azure.AI.Agents.Persistent" Version="1.2.0-beta.1" />
<PackageVersion Include="Azure.AI.OpenAI" Version="2.3.0-beta.1" />
<PackageVersion Include="Azure.Identity" Version="1.14.2" />
@@ -61,6 +65,7 @@
<PackageVersion Include="Microsoft.Extensions.Logging.Abstractions" Version="9.0.8" />
<PackageVersion Include="Microsoft.Extensions.Logging.Console" Version="9.0.8" />
<PackageVersion Include="Microsoft.Extensions.Logging.Testing" Version="9.0.8" />
<PackageVersion Include="Microsoft.Extensions.Options" Version="9.0.8" />
<!-- Vector Stores -->
<PackageVersion Include="Microsoft.SemanticKernel.Connectors.InMemory" Version="1.61.0-preview" />
<!-- Agent SDKs -->
+5
View File
@@ -28,6 +28,7 @@
<Folder Name="/Solution Items/.github/workflows/">
<File Path="../.github/workflows/dotnet-build-and-test.yml" />
<File Path="../.github/workflows/dotnet-check-coverage.ps1" />
<File Path="../.github/workflows/dotnet-cosmosdb-integration-tests.yml" />
<File Path="../.github/workflows/dotnet-format.yml" />
</Folder>
<Folder Name="/Solution Items/docs/" />
@@ -121,6 +122,10 @@
<Project Path="src/Microsoft.Extensions.AI.Agents/Microsoft.Extensions.AI.Agents.csproj" />
</Folder>
<Folder Name="/Tests/" />
<Folder Name="/Tests/Cosmos/">
<Project Path="tests/CosmosDB.IntegrationTests/CosmosDB.Testing.AppHost/CosmosDB.Testing.AppHost.csproj" />
<Project Path="tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj" />
</Folder>
<Folder Name="/Tests/IntegrationTests/">
<Project Path="tests/AgentConformance.IntegrationTests/AgentConformance.IntegrationTests.csproj" />
<Project Path="tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistent.IntegrationTests.csproj" />
@@ -1,6 +1,6 @@
<Project Sdk="Microsoft.NET.Sdk">
<Sdk Name="Aspire.AppHost.Sdk" Version="9.3.1" />
<Sdk Name="Aspire.AppHost.Sdk" Version="$(AspireAppHostSdkVersion)" />
<PropertyGroup>
<OutputType>Exe</OutputType>
@@ -11,13 +11,6 @@ namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
/// the entire actor's state for optimistic concurrency control.
/// This document contains no actor state data. It only serves to track last modified
/// time and provide a single ETag for the actor's state.
///
/// Example structure:
/// {
/// "id": "rootdoc", // Root document ID (constant per actor partition)
/// "actorId": "actor-123", // Partition key (actor ID)
/// "lastModified": "2024-...", // Timestamp
/// }
/// </summary>
public sealed class ActorRootDocument
{
@@ -27,9 +20,13 @@ public sealed class ActorRootDocument
public string Id { get; set; } = default!;
/// <summary>
/// The actor ID.
/// The actor type.
/// </summary>
public string ActorId { get; set; } = default!;
public string ActorType { get; set; } = default!;
/// <summary>
/// The actor key.
/// </summary>
public string ActorKey { get; set; } = default!;
/// <summary>
/// The last modified timestamp.
@@ -55,9 +52,13 @@ public sealed class ActorStateDocument
public string Id { get; set; } = default!;
/// <summary>
/// The actor ID.
/// The actor type.
/// </summary>
public string ActorId { get; set; } = default!;
public string ActorType { get; set; } = default!;
/// <summary>
/// The actor key.
/// </summary>
public string ActorKey { get; set; } = default!;
/// <summary>
/// The logical key for the state entry.
@@ -11,7 +11,7 @@ namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
/// <summary>
/// Cosmos DB implementation of actor state storage.
/// </summary>
public class CosmosActorStateStorage : IActorStateStorage
public class CosmosActorStateStorage : IActorStateStorage, IAsyncDisposable
{
private readonly LazyCosmosContainer _lazyContainer;
private const string InitialEtag = "0"; // Initial ETag value when no state exists
@@ -47,8 +47,8 @@ public class CosmosActorStateStorage : IActorStateStorage
}
var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false);
var batch = container.CreateTransactionalBatch(GetPartitionKey(actorId));
var actorIdStr = actorId.ToString();
var (partitionKey, actorType, actorKey) = BuildPartitionKey(actorId);
var batch = container.CreateTransactionalBatch(partitionKey);
// Add data operations to batch
foreach (var op in operations)
@@ -61,7 +61,8 @@ public class CosmosActorStateStorage : IActorStateStorage
var item = new ActorStateDocument
{
Id = docId,
ActorId = actorIdStr,
ActorType = actorType,
ActorKey = actorKey,
Key = set.Key,
Value = set.Value
};
@@ -83,7 +84,8 @@ public class CosmosActorStateStorage : IActorStateStorage
var newRoot = new ActorRootDocument
{
Id = RootDocumentId,
ActorId = actorId.ToString(),
ActorType = actorType,
ActorKey = actorKey,
LastModified = DateTimeOffset.UtcNow,
};
@@ -103,6 +105,7 @@ public class CosmosActorStateStorage : IActorStateStorage
var result = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false);
if (!result.IsSuccessStatusCode)
{
_ = result.ErrorMessage;
return new WriteResponse(eTag: string.Empty, success: false);
}
@@ -135,6 +138,8 @@ public class CosmosActorStateStorage : IActorStateStorage
// Read root document first to get actor-level ETag
string actorETag = await this.GetActorETagAsync(container, actorId, cancellationToken).ConfigureAwait(false);
var actorType = actorId.Type.ToString();
var actorKey = actorId.Key;
foreach (var op in operations)
{
@@ -162,14 +167,16 @@ public class CosmosActorStateStorage : IActorStateStorage
QueryDefinition query;
if (!string.IsNullOrEmpty(list.KeyPrefix))
{
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorId = @actorId AND c.key != null AND STARTSWITH(c.key, @keyPrefix)")
.WithParameter("@actorId", actorId.ToString())
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorType = @actorType AND c.actorKey = @actorKey AND c.key != null AND STARTSWITH(c.key, @keyPrefix)")
.WithParameter("@actorType", actorType)
.WithParameter("@actorKey", actorKey)
.WithParameter("@keyPrefix", list.KeyPrefix);
}
else
{
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorId = @actorId AND c.key != null")
.WithParameter("@actorId", actorId.ToString());
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorType = @actorType AND c.actorKey = @actorKey AND c.key != null")
.WithParameter("@actorType", actorType)
.WithParameter("@actorKey", actorKey);
}
var requestOptions = new QueryRequestOptions
@@ -212,7 +219,18 @@ public class CosmosActorStateStorage : IActorStateStorage
private const string RootDocumentId = "rootdoc";
private static PartitionKey GetPartitionKey(ActorId actorId)
=> new(actorId.ToString());
{
var (partitionKey, _, _) = BuildPartitionKey(actorId);
return partitionKey;
}
private static (PartitionKey partitionKey, string actorType, string actorKey) BuildPartitionKey(ActorId actorId)
{
var actorType = actorId.Type.ToString();
var actorKey = actorId.Key;
var partitionKey = new PartitionKeyBuilder().Add(actorType).Add(actorKey).Build();
return (partitionKey, actorType, actorKey);
}
/// <summary>
/// Gets the current ETag for the actor's root document.
@@ -234,4 +252,13 @@ public class CosmosActorStateStorage : IActorStateStorage
return InitialEtag;
}
}
/// <summary>
/// Disposes the Cosmos DB container asynchronously.
/// </summary>
public async ValueTask DisposeAsync()
{
await this._lazyContainer.DisposeAsync().ConfigureAwait(false);
GC.SuppressFinalize(this);
}
}
@@ -1,26 +1,30 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.ObjectModel;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Azure.Cosmos;
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
#pragma warning disable VSTHRD011 // Use AsyncLazy<T>
/// <summary>
/// A lazy wrapper around a Cosmos DB Container.
/// This avoids performing async I/O-bound operations (i.e. Cosmos DB setup) during
/// DI registration, deferring them until first access.
/// </summary>
internal sealed class LazyCosmosContainer
internal sealed class LazyCosmosContainer : IAsyncDisposable
{
private readonly static Random s_random = new();
private readonly CosmosClient? _cosmosClient;
private readonly string? _databaseName;
private readonly string? _containerName;
private readonly Lazy<Task<Container>> _lazyContainer;
private readonly CancellationTokenSource _cts = new();
private Task<Container>? _initTask;
// internal for testing
internal readonly static string[] CosmosPartitionKeyPaths = ["/actorType", "/actorKey"];
/// <summary>
/// LazyCosmosContainer constructor that initializes the container lazily.
@@ -30,9 +34,7 @@ internal sealed class LazyCosmosContainer
this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient));
this._databaseName = databaseName ?? throw new ArgumentNullException(nameof(databaseName));
this._containerName = containerName ?? throw new ArgumentNullException(nameof(containerName));
this._lazyContainer = new Lazy<Task<Container>>(this.InitializeContainerAsync, LazyThreadSafetyMode.ExecutionAndPublication);
}
/// <summary>
/// LazyCosmosContainer constructor that accepts an existing Container instance.
/// </summary>
@@ -43,20 +45,59 @@ internal sealed class LazyCosmosContainer
throw new ArgumentNullException(nameof(container));
}
this._lazyContainer = new Lazy<Task<Container>>(() => Task.FromResult(container), LazyThreadSafetyMode.ExecutionAndPublication);
this._initTask = Task.FromResult(container);
}
/// <summary>
/// Gets the Container, initializing it if necessary.
/// </summary>
public Task<Container> GetContainerAsync() => this._lazyContainer.Value;
public Task<Container> GetContainerAsync()
=> this._initTask ??= this.InitializeWithRetryAsync(this._cts.Token);
private async Task<Container> InitializeContainerAsync()
private async Task<Container> InitializeWithRetryAsync(CancellationToken cancellationToken)
{
var baseDelay = TimeSpan.FromSeconds(1);
var maxDelay = TimeSpan.FromSeconds(30);
var previousDelay = baseDelay;
while (true)
{
cancellationToken.ThrowIfCancellationRequested();
try
{
return await this.InitializeContainerAsync(cancellationToken).ConfigureAwait(false);
}
catch (CosmosException ex) when (IsTransient(ex))
{
// If server provided RetryAfter, respect it but add a small jitter so clients don't retry in perfect sync.
if (ex.RetryAfter is not null && ex.RetryAfter > TimeSpan.Zero)
{
var retry = ex.RetryAfter.Value;
var jitterMs = this.RandomNextDouble() * retry.TotalMilliseconds; // 0..retry
var delay = retry + TimeSpan.FromMilliseconds(jitterMs);
await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
previousDelay = delay;
continue;
}
// sleep = min(maxDelay, random(baseDelay, previousDelay * 3))
var minMs = baseDelay.TotalMilliseconds;
var maxMs = Math.Min(maxDelay.TotalMilliseconds, Math.Max(minMs, previousDelay.TotalMilliseconds * 3));
var sleepMs = this.RandomNextDouble() * (maxMs - minMs) + minMs;
var jitterDelay = TimeSpan.FromMilliseconds(sleepMs);
await Task.Delay(jitterDelay, cancellationToken).ConfigureAwait(false);
previousDelay = jitterDelay;
}
}
}
private async Task<Container> InitializeContainerAsync(CancellationToken cancellationToken)
{
// Create database if it doesn't exist
var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!).ConfigureAwait(false);
var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!, cancellationToken: cancellationToken).ConfigureAwait(false);
var containerProperties = new ContainerProperties(this._containerName!, "/actorId")
var containerProperties = new ContainerProperties(this._containerName!, CosmosPartitionKeyPaths)
{
Id = this._containerName!,
IndexingPolicy = new IndexingPolicy
@@ -64,17 +105,50 @@ internal sealed class LazyCosmosContainer
IndexingMode = IndexingMode.Consistent,
Automatic = true
},
PartitionKeyPaths = ["/actorId"]
PartitionKeyPaths = CosmosPartitionKeyPaths
};
// Add composite index for efficient queries
containerProperties.IndexingPolicy.CompositeIndexes.Add(new Collection<CompositePath>
{
new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending },
new() { Path = "/actorType", Order = CompositePathSortOrder.Ascending },
new() { Path = "/actorKey", Order = CompositePathSortOrder.Ascending },
new() { Path = "/key", Order = CompositePathSortOrder.Ascending }
});
var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties).ConfigureAwait(false);
var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: cancellationToken).ConfigureAwait(false);
return container.Container;
}
private static bool IsTransient(Exception exception)
{
return exception switch
{
CosmosException cosmosEx => cosmosEx.StatusCode switch
{
#if NET9_0_OR_GREATER
HttpStatusCode.TooManyRequests => true, // 429 - Rate limited
#endif
HttpStatusCode.InternalServerError => true, // 500 - Server error
HttpStatusCode.BadGateway => true, // 502 - Bad gateway
HttpStatusCode.ServiceUnavailable => true, // 503 - Service unavailable
HttpStatusCode.GatewayTimeout => true, // 504 - Gateway timeout
HttpStatusCode.RequestTimeout => true, // 408 - Request timeout
_ => false
},
TaskCanceledException or OperationCanceledException or ArgumentException => false,
_ => true // Retry other exceptions (network issues, etc.)
};
}
#pragma warning disable CA5394 // Do not use insecure randomness
private double RandomNextDouble() => s_random.NextDouble();
#pragma warning restore CA5394 // Do not use insecure randomness
public ValueTask DisposeAsync()
{
this._cts?.Cancel();
this._cts?.Dispose();
return default;
}
}
@@ -9,6 +9,7 @@
<ItemGroup>
<PackageReference Include="Microsoft.Azure.Cosmos" />
<PackageReference Include="Microsoft.Extensions.Options" />
<PackageReference Include="Newtonsoft.Json" />
</ItemGroup>
@@ -1,10 +1,18 @@
// Copyright (c) Microsoft. All rights reserved.
using Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests;
using CosmosDB.Testing.AppHost;
var builder = DistributedApplication.CreateBuilder(args);
var cosmosDb = builder.AddAzureCosmosDB(CosmosDBTestConstants.TestCosmosDbName);
if (CosmosDBTestConstants.UseEmulatorForTesting)
if (CosmosDBTestConstants.UseEmulatorInCICD)
{
// Emulator created in the CI/CD pipeline gives more control over some settings and port-configuration today.
// It probably should be configured here to use 8081 port + setup the partition count and throughput, but it's not supported in Aspire yet, so leaving as a placeholder.
// Once Aspire's emulator is suported, the emulator in CI/CD can be removed.
cosmosDb.RunAsEmulator(emulator => emulator.WithLifetime(ContainerLifetime.Persistent));
}
else if (CosmosDBTestConstants.UseAspireEmulatorForTesting)
{
cosmosDb.RunAsEmulator(emulator => emulator.WithLifetime(ContainerLifetime.Persistent));
}
@@ -1,16 +1,15 @@
<Project Sdk="Microsoft.NET.Sdk">
<Project Sdk="Microsoft.NET.Sdk">
<Sdk Name="Aspire.AppHost.Sdk" Version="9.3.1" />
<Sdk Name="Aspire.AppHost.Sdk" Version="$(AspireAppHostSdkVersion)" />
<PropertyGroup>
<!--<OutputType>Exe</OutputType>-->
<OutputType>Exe</OutputType>
<TargetFramework>net9.0</TargetFramework>
<ImplicitUsings>enable</ImplicitUsings>
<Nullable>enable</Nullable>
<UserSecretsId>80048040-aaf1-4f44-9970-8aef39651edf</UserSecretsId>
<IsAspireHost>true</IsAspireHost>
<UserSecretsId>cb8630a8-ec5e-4676-a2b0-4497965c809d</UserSecretsId>
<GenerateProgramFile>false</GenerateProgramFile>
</PropertyGroup>
</PropertyGroup>
<ItemGroup>
<PackageReference Include="Aspire.Hosting.AppHost" />
@@ -0,0 +1,24 @@
// Copyright (c) Microsoft. All rights reserved.
//using System.Linq.Expressions;
namespace CosmosDB.Testing.AppHost;
public static class CosmosDBTestConstants
{
public const string TestCosmosDbName = "ActorStateStorageTests";
public const string TestCosmosDbDatabaseName = "state-database";
//Set to use the CosmosDB emulator for testing via environment variable.
//Example: set COSMOSDB_TESTS_USE_EMULATOR = true in your environment.
//Warning: Using the emulator may cause test flakiness.
public static bool UseAspireEmulatorForTesting => string.Equals(
Environment.GetEnvironmentVariable("COSMOSDB_TESTS_USE_EMULATOR"),
"true",
StringComparison.OrdinalIgnoreCase);
public static bool UseEmulatorInCICD => string.Equals(
Environment.GetEnvironmentVariable("COSMOSDB_TESTS_USE_EMULATOR_CICD"),
"true",
StringComparison.OrdinalIgnoreCase);
}
@@ -21,6 +21,7 @@
"environmentVariables": {
"ASPNETCORE_ENVIRONMENT": "Development",
"DOTNET_ENVIRONMENT": "Development",
"ASPIRE_ALLOW_UNSECURED_TRANSPORT": "true",
"ASPIRE_DASHBOARD_OTLP_ENDPOINT_URL": "http://localhost:19080",
"ASPIRE_RESOURCE_SERVICE_ENDPOINT_URL": "http://localhost:20201"
}
@@ -0,0 +1,8 @@
{
"Logging": {
"LogLevel": {
"Default": "Information",
"Microsoft.AspNetCore": "Warning"
}
}
}
@@ -0,0 +1,9 @@
{
"Logging": {
"LogLevel": {
"Default": "Information",
"Microsoft.AspNetCore": "Warning",
"Aspire.Hosting.Dcp": "Warning"
}
}
}
@@ -1,18 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests;
public static class CosmosDBTestConstants
{
public const string TestCosmosDbName = "ActorStateStorageTests";
public const string TestCosmosDbDatabaseName = "state-database";
// Set to use the CosmosDB emulator for testing via environment variable.
// Example: set COSMOSDB_TESTS_USE_EMULATOR=true in your environment.
// Warning: Using the emulator may cause test flakiness.
public static bool UseEmulatorForTesting =>
string.Equals(
Environment.GetEnvironmentVariable("COSMOSDB_TESTS_USE_EMULATOR"),
"true",
StringComparison.OrdinalIgnoreCase);
}
@@ -31,7 +31,7 @@ public class CosmosActorStateStorageConcurrencyTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key = "testKey";
@@ -70,7 +70,7 @@ public class CosmosActorStateStorageConcurrencyTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
// Setup initial state
@@ -185,7 +185,7 @@ public class CosmosActorStateStorageConcurrencyTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key = "testKey";
@@ -258,7 +258,7 @@ public class CosmosActorStateStorageConcurrencyTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); // Fresh actor
var key = "testKey";
@@ -308,7 +308,7 @@ public class CosmosActorStateStorageConcurrencyTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); // Non-existent actor
var key = "testKey";
@@ -26,7 +26,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var prefixKey1 = "prefix_key1";
@@ -69,7 +69,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key1 = "key1";
@@ -107,7 +107,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
// Act - List keys for actor with no state
@@ -132,7 +132,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key1 = "key1";
@@ -173,7 +173,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key1 = "key1";
@@ -225,7 +225,7 @@ public class CosmosActorStateStorageListKeysTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
// Create keys with different prefixes
@@ -26,7 +26,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key = "testKey";
@@ -52,7 +52,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key1 = "key1";
@@ -154,7 +154,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key = "testKey";
@@ -196,7 +196,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId1 = new ActorId("TestActor1", Guid.NewGuid().ToString());
var testActorId2 = new ActorId("TestActor2", Guid.NewGuid().ToString());
@@ -242,7 +242,7 @@ public class CosmosActorStateStorageTests
// Arrange
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var emptyOperations = new List<ActorStateWriteOperation>();
// Act & Assert
@@ -259,7 +259,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var readOperations = new List<ActorStateReadOperation>
@@ -284,7 +284,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
// Create a complex object with various types
@@ -357,7 +357,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key1 = "key1";
@@ -424,7 +424,7 @@ public class CosmosActorStateStorageTests
using var cts = new CancellationTokenSource(s_defaultTimeout);
var cancellationToken = cts.Token;
var storage = new CosmosActorStateStorage(this._fixture.Container);
await using var storage = new CosmosActorStateStorage(this._fixture.Container);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
// Test keys with special characters that need sanitization
@@ -3,6 +3,7 @@
using System.Text.Json;
using Aspire.Hosting;
using Azure.Identity;
using CosmosDB.Testing.AppHost;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.Logging;
@@ -30,7 +31,7 @@ public class CosmosTestFixture : IAsyncLifetime
var cancellationToken = cts.Token;
var appHost = await DistributedApplicationTestingBuilder
.CreateAsync<Projects.Microsoft_Extensions_AI_Agents_Runtime_Storage_CosmosDB_Tests_AppHost>(cancellationToken);
.CreateAsync<Projects.CosmosDB_Testing_AppHost>(cancellationToken);
appHost.Services.AddLogging(logging =>
{
@@ -47,7 +48,16 @@ public class CosmosTestFixture : IAsyncLifetime
this.App = await appHost.BuildAsync(cancellationToken).WaitAsync(cancellationToken);
await this.App.StartAsync(cancellationToken).WaitAsync(cancellationToken);
var cs = await this.App.GetConnectionStringAsync(CosmosDBTestConstants.TestCosmosDbName, cancellationToken);
var connectionString = await this.App.GetConnectionStringAsync(CosmosDBTestConstants.TestCosmosDbName, cancellationToken);
if (CosmosDBTestConstants.UseEmulatorInCICD)
{
// Emulator is setup in the CI/CD pipeline, so we will not use one produced by Aspire.
// For simplicity, we override the connection string here with the well-known emulator connection string.
// https://learn.microsoft.com/en-us/azure/cosmos-db/emulator
connectionString = "AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==;";
}
CosmosClientOptions ccoptions = new()
{
UseSystemTextJsonSerializerWithOptions = new JsonSerializerOptions()
@@ -57,23 +67,29 @@ public class CosmosTestFixture : IAsyncLifetime
}
};
if (CosmosDBTestConstants.UseEmulatorForTesting)
if (CosmosDBTestConstants.UseAspireEmulatorForTesting || CosmosDBTestConstants.UseEmulatorInCICD)
{
ccoptions.ConnectionMode = ConnectionMode.Gateway;
ccoptions.LimitToEndpoint = true;
this.CosmosClient = new CosmosClient(cs, ccoptions);
this.CosmosClient = new CosmosClient(connectionString, ccoptions);
}
else
{
this.CosmosClient = new CosmosClient(cs, new DefaultAzureCredential(), ccoptions);
this.CosmosClient = new CosmosClient(connectionString, new DefaultAzureCredential(), ccoptions);
}
var database = this.CosmosClient.GetDatabase(CosmosDBTestConstants.TestCosmosDbDatabaseName);
// raise throughput to avoid parallel test execution failures
var throughputProperties = ThroughputProperties.CreateAutoscaleThroughput(100000);
// Ensure database exists. It will be a no-op if it was already created before.
_ = await this.CosmosClient.CreateDatabaseIfNotExistsAsync(CosmosDBTestConstants.TestCosmosDbDatabaseName, throughputProperties);
var containerProperties = new ContainerProperties()
{
Id = "CosmosActorStateStorageTests",
PartitionKeyPath = "/actorId"
PartitionKeyPaths = LazyCosmosContainer.CosmosPartitionKeyPaths
};
this.Container = await database.CreateContainerIfNotExistsAsync(containerProperties);
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Text.Json;
using CosmosDB.Testing.AppHost;
using Microsoft.Azure.Cosmos;
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests;
@@ -25,7 +26,7 @@ public class LazyCosmosContainerTests
{
// Arrange
using var cts = new CancellationTokenSource(s_defaultTimeout);
var lazyContainer = new LazyCosmosContainer(this._fixture.Container);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.Container);
// Act
var result = await lazyContainer.GetContainerAsync();
@@ -39,7 +40,7 @@ public class LazyCosmosContainerTests
{
// Arrange
using var cts = new CancellationTokenSource(s_defaultTimeout);
var lazyContainer = new LazyCosmosContainer(this._fixture.Container);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.Container);
// Act
var result1 = await lazyContainer.GetContainerAsync();
@@ -52,7 +53,7 @@ public class LazyCosmosContainerTests
Assert.Same(this._fixture.Container, result1);
}
[Fact]
[SkipOnEmulatorFact]
public async Task GetContainerAsync_WithCosmosClient_ShouldInitializeAndWorkCorrectlyAsync()
{
// Arrange
@@ -61,7 +62,7 @@ public class LazyCosmosContainerTests
// Create a unique container name for this test
var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}";
var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
try
{
@@ -74,7 +75,7 @@ public class LazyCosmosContainerTests
// Verify the container can perform basic operations
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var storage = new CosmosActorStateStorage(lazyContainer);
await using var storage = new CosmosActorStateStorage(lazyContainer);
var key = "testKey";
var value = JsonSerializer.SerializeToElement("testValue");
@@ -111,7 +112,7 @@ public class LazyCosmosContainerTests
// Create a unique container name for this test
var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}";
var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
try
{
@@ -148,7 +149,7 @@ public class LazyCosmosContainerTests
// Create a unique container name for this test
var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}";
var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
try
{
@@ -211,7 +212,7 @@ public class LazyCosmosContainerTests
Assert.Throws<ArgumentNullException>(() => new LazyCosmosContainer(this._fixture.CosmosClient, "test-db", null!));
}
[Fact]
[SkipOnEmulatorFact]
public async Task LazyCosmosContainer_WithInternalConstructor_ShouldWorkWithCosmosActorStateStorageAsync()
{
// Arrange
@@ -220,12 +221,12 @@ public class LazyCosmosContainerTests
// Create a unique container name for this test
var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}";
var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
await using var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName);
try
{
// Act - Create storage using the internal constructor (like DI would)
var storage = new CosmosActorStateStorage(lazyContainer);
await using var storage = new CosmosActorStateStorage(lazyContainer);
var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString());
var key = "testKey";
@@ -277,7 +278,7 @@ public class LazyCosmosContainerTests
// Use an invalid database name that should cause Cosmos to reject it
var invalidDatabaseName = new string('a', 256); // Database names have limits
var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, invalidDatabaseName, "test-container");
await using var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, invalidDatabaseName, "test-container");
// Act & Assert
await Assert.ThrowsAsync<CosmosException>(async () => await lazyContainer.GetContainerAsync());
@@ -13,8 +13,8 @@
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj" />
<ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj" />
<ProjectReference Include="..\CosmosDB.Testing.AppHost\CosmosDB.Testing.AppHost.csproj" />
</ItemGroup>
<ItemGroup>
@@ -1,5 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using CosmosDB.Testing.AppHost;
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests;
/// <summary>
@@ -13,9 +15,14 @@ public sealed class SkipOnEmulatorFactAttribute : FactAttribute
/// </summary>
public SkipOnEmulatorFactAttribute()
{
if (CosmosDBTestConstants.UseEmulatorForTesting)
if (CosmosDBTestConstants.UseAspireEmulatorForTesting)
{
this.Skip = "Skipping test on CosmosDB emulator.";
this.Skip = "Skipping test on Aspire-configured CosmosDB emulator.";
}
if (CosmosDBTestConstants.UseEmulatorInCICD)
{
this.Skip = "Skipping test on CICD-configured CosmosDB emulator.";
}
}
}