mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
.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:
committed by
GitHub
Unverified
parent
1116916737
commit
25291de8cb
@@ -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
|
||||
@@ -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 -->
|
||||
|
||||
@@ -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>
|
||||
|
||||
+12
-11
@@ -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.
|
||||
|
||||
+37
-10
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+90
-16
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
+1
@@ -9,6 +9,7 @@
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="Microsoft.Azure.Cosmos" />
|
||||
<PackageReference Include="Microsoft.Extensions.Options" />
|
||||
<PackageReference Include="Newtonsoft.Json" />
|
||||
</ItemGroup>
|
||||
|
||||
|
||||
+10
-2
@@ -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));
|
||||
}
|
||||
+5
-6
@@ -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" />
|
||||
+24
@@ -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);
|
||||
}
|
||||
+1
@@ -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"
|
||||
}
|
||||
+8
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
-18
@@ -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);
|
||||
}
|
||||
+5
-5
@@ -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";
|
||||
|
||||
+6
-6
@@ -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
|
||||
|
||||
+9
-9
@@ -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
|
||||
|
||||
+22
-6
@@ -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);
|
||||
|
||||
+12
-11
@@ -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());
|
||||
|
||||
+1
-1
@@ -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>
|
||||
|
||||
+9
-2
@@ -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.";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user