mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
.NET: Add Cosmos DB implementations for ChatMessageStore and CheckpointStore. (#1838)
* draft commit * Added Cosmos agent thread and tests * revert unnecessary changes and fix tests * add multi-tenant support with hierarchical partition keys (and tests). * enhance transactional batch * address review comments * Address PR review comments from @westey-m * Merge upstream/main - resolve slnx conflicts * use param validation helpers * Replace useManagedIdentity boolean with TokenCredential parameter * Remove redundant suppressions and fix tests * Rename project from Microsoft.Agents.AI.Abstractions.CosmosNoSql to Microsoft.Agents.AI.CosmosNoSql * Refactor constructors to use chaining pattern * Reorder deserialization constructor parameters for consistency * Remove database/container IDs from serialized state * Remove auto-generation of MessageId * Optimize AddMessagesAsync to avoid enumeration when possible * Add MaxMessagesToRetrieve to limit context window * Make Role nullable instead of defaulting * Fix net472 build without rebasing 19 commits * Add Cosmos DB emulator to CI workflow * Fix Cosmos DB emulator tests: use Skip.If instead of Assert.Fail and start emulator before unit tests * Replace Skip.If() with conditional return to fix compilation * Use env var to skip Cosmos tests on non-Windows CI * Add Xunit.SkippableFact package to properly skip Cosmos tests on Linux * Change [Fact] to [SkippableFact] for proper test skipping behavior * Remove stale Microsoft.Agents.AI.Abstractions.CosmosNoSql directory * Fix code formatting: add braces, this. qualifications, and final newlines * Fix file encoding to UTF-8 with BOM, fix import ordering, and remove unnecessary using directives * Convert backing fields to auto-properties and remove Azure.Identity using directive * Fix CosmosChatMessageStore.cs encoding back to UTF-8 with BOM * Fix test file formatting: indentation, encoding, imports, this. qualifications, naming conventions, and simplify new expressions * Fix const field naming violations: Remove s_ prefix from const fields and add this. qualification to Dispose call * Add local .editorconfig for Cosmos DB tests to suppress IDE0005 false positives from multi-targeting * Fix IDE1006 naming violations: Rename TestDatabaseId to s_testDatabaseId and add final newlines * Address PR review comments Address Wesley's review comments: - Remove Cosmos DB package references from core projects - Delete duplicate test files from old package structure - Remove redundant parameter validation from extension methods Address Kiran's review comments: - Remove redundant 429 retry logic (SDK handles automatically) - Add explicit RequestEntityTooLarge error handling - Remove dead code in GetMessageCountAsync - Add defensive partition key validation comments * Fix IDE0001 formatting error in AgentProviderExtensions.cs. Use type alias to resolve namespace conflict between Azure.AI.Agents.Persistent.RunStatus and Microsoft.Agents.AI.Workflows.RunStatus. This eliminates the need for global:: qualifier which triggered the formatter warning. * Update package versions for Aspire 13.0.0 compatibility * Fix TargetFrameworks in Cosmos DB projects - Replace with which is defined in Directory.Build.props - Fix package reference from System.Linq.Async to System.Linq.AsyncEnumerable to match Directory.Packages.props * Remove redundant counter, add partition key validation, use factory pattern for deserialization
This commit is contained in:
Theo van Kraay
committed by
GitHub
Unverified
GitHub
parent
907d79ab3c
commit
a57b37d5fa
@@ -0,0 +1,9 @@
|
||||
# EditorConfig overrides for Cosmos DB Unit Tests
|
||||
# Multi-targeting (net472 + net9.0) causes false positives for IDE0005 (unnecessary using directives)
|
||||
|
||||
root = false
|
||||
|
||||
[*.cs]
|
||||
# Suppress IDE0005 for this project - multi-targeting causes false positives
|
||||
# These using directives ARE necessary but appear unnecessary in one target framework
|
||||
dotnet_diagnostic.IDE0005.severity = none
|
||||
+760
@@ -0,0 +1,760 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Linq;
|
||||
using System.Text.Json;
|
||||
using System.Text.Json.Serialization.Metadata;
|
||||
using System.Threading.Tasks;
|
||||
using Azure.Core;
|
||||
using Azure.Identity;
|
||||
using Microsoft.Agents.AI;
|
||||
using Microsoft.Azure.Cosmos;
|
||||
using Microsoft.Extensions.AI;
|
||||
using Xunit;
|
||||
|
||||
namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests;
|
||||
|
||||
/// <summary>
|
||||
/// Contains tests for <see cref="CosmosChatMessageStore"/>.
|
||||
///
|
||||
/// Test Modes:
|
||||
/// - Default Mode: Cleans up all test data after each test run (deletes database)
|
||||
/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer
|
||||
///
|
||||
/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true
|
||||
/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test
|
||||
///
|
||||
/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at:
|
||||
/// https://localhost:8081/_explorer/index.html
|
||||
/// Database: AgentFrameworkTests
|
||||
/// Container: ChatMessages
|
||||
///
|
||||
/// Environment Variable Reference:
|
||||
/// | Variable | Values | Description |
|
||||
/// |----------|--------|-------------|
|
||||
/// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion |
|
||||
///
|
||||
/// Usage Examples:
|
||||
/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/
|
||||
/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB"
|
||||
/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/
|
||||
/// </summary>
|
||||
[Collection("CosmosDB")]
|
||||
public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable
|
||||
{
|
||||
// Cosmos DB Emulator connection settings
|
||||
private const string EmulatorEndpoint = "https://localhost:8081";
|
||||
private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==";
|
||||
private const string TestContainerId = "ChatMessages";
|
||||
private const string HierarchicalTestContainerId = "HierarchicalChatMessages";
|
||||
// Use unique database ID per test class instance to avoid conflicts
|
||||
#pragma warning disable CA1802 // Use literals where appropriate
|
||||
private static readonly string s_testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}";
|
||||
#pragma warning restore CA1802
|
||||
|
||||
private string _connectionString = string.Empty;
|
||||
private bool _emulatorAvailable;
|
||||
private bool _preserveContainer;
|
||||
private CosmosClient? _setupClient; // Only used for test setup/cleanup
|
||||
|
||||
public async Task InitializeAsync()
|
||||
{
|
||||
// Check environment variable to determine if we should preserve containers
|
||||
// Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection
|
||||
this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}";
|
||||
|
||||
try
|
||||
{
|
||||
// Only create CosmosClient for test setup - the actual tests will use connection string constructors
|
||||
this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey);
|
||||
|
||||
// Test connection by attempting to create database
|
||||
var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId);
|
||||
|
||||
// Create container for simple partitioning tests
|
||||
await databaseResponse.Database.CreateContainerIfNotExistsAsync(
|
||||
TestContainerId,
|
||||
"/conversationId",
|
||||
throughput: 400);
|
||||
|
||||
// Create container for hierarchical partitioning tests with hierarchical partition key
|
||||
var hierarchicalContainerProperties = new ContainerProperties(HierarchicalTestContainerId, new List<string> { "/tenantId", "/userId", "/sessionId" });
|
||||
await databaseResponse.Database.CreateContainerIfNotExistsAsync(
|
||||
hierarchicalContainerProperties,
|
||||
throughput: 400);
|
||||
|
||||
this._emulatorAvailable = true;
|
||||
}
|
||||
catch (Exception)
|
||||
{
|
||||
// Emulator not available, tests will be skipped
|
||||
this._emulatorAvailable = false;
|
||||
this._setupClient?.Dispose();
|
||||
this._setupClient = null;
|
||||
}
|
||||
}
|
||||
|
||||
public async Task DisposeAsync()
|
||||
{
|
||||
if (this._setupClient != null && this._emulatorAvailable)
|
||||
{
|
||||
try
|
||||
{
|
||||
if (this._preserveContainer)
|
||||
{
|
||||
// Preserve mode: Don't delete the database/container, keep data for inspection
|
||||
// This allows viewing data in the Cosmos DB Emulator Data Explorer
|
||||
// No cleanup needed - data persists for debugging
|
||||
}
|
||||
else
|
||||
{
|
||||
// Clean mode: Delete the test database and all data
|
||||
var database = this._setupClient.GetDatabase(s_testDatabaseId);
|
||||
await database.DeleteAsync();
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
// Ignore cleanup errors during test teardown
|
||||
Console.WriteLine($"Warning: Cleanup failed: {ex.Message}");
|
||||
}
|
||||
finally
|
||||
{
|
||||
this._setupClient.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public void Dispose()
|
||||
{
|
||||
this._setupClient?.Dispose();
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
|
||||
private void SkipIfEmulatorNotAvailable()
|
||||
{
|
||||
// In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true"
|
||||
// Locally: Skip if emulator connection check failed
|
||||
var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available");
|
||||
}
|
||||
|
||||
#region Constructor Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithConnectionString_ShouldCreateInstance()
|
||||
{
|
||||
// Arrange & Act
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, "test-conversation");
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(store);
|
||||
Assert.Equal("test-conversation", store.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(TestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(store);
|
||||
Assert.NotNull(store.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(TestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithNullConnectionString_ShouldThrowArgumentException()
|
||||
{
|
||||
// Arrange & Act & Assert
|
||||
Assert.Throws<ArgumentNullException>(() =>
|
||||
new CosmosChatMessageStore((string)null!, s_testDatabaseId, TestContainerId, "test-conversation"));
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException()
|
||||
{
|
||||
// Arrange & Act & Assert
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
Assert.Throws<ArgumentException>(() =>
|
||||
new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ""));
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region AddMessagesAsync Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var conversationId = Guid.NewGuid().ToString();
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId);
|
||||
var message = new ChatMessage(ChatRole.User, "Hello, world!");
|
||||
|
||||
// Act
|
||||
await store.AddMessagesAsync([message]);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Assert
|
||||
var messages = await store.GetMessagesAsync();
|
||||
var messageList = messages.ToList();
|
||||
|
||||
// Simple assertion - if this fails, we know the deserialization is the issue
|
||||
if (messageList.Count == 0)
|
||||
{
|
||||
// Let's check if we can find ANY items in the container for this conversation
|
||||
var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId")
|
||||
.WithParameter("@conversationId", conversationId);
|
||||
var countIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId)
|
||||
.GetItemQueryIterator<int>(directQuery, requestOptions: new QueryRequestOptions
|
||||
{
|
||||
PartitionKey = new PartitionKey(conversationId)
|
||||
});
|
||||
|
||||
var countResponse = await countIterator.ReadNextAsync();
|
||||
var count = countResponse.FirstOrDefault();
|
||||
|
||||
// Debug: Let's see what the raw query returns
|
||||
var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId")
|
||||
.WithParameter("@conversationId", conversationId);
|
||||
var rawIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId)
|
||||
.GetItemQueryIterator<dynamic>(rawQuery, requestOptions: new QueryRequestOptions
|
||||
{
|
||||
PartitionKey = new PartitionKey(conversationId)
|
||||
});
|
||||
|
||||
List<dynamic> rawResults = new();
|
||||
while (rawIterator.HasMoreResults)
|
||||
{
|
||||
var rawResponse = await rawIterator.ReadNextAsync();
|
||||
rawResults.AddRange(rawResponse);
|
||||
}
|
||||
|
||||
string rawJson = rawResults.Count > 0 ? Newtonsoft.Json.JsonConvert.SerializeObject(rawResults[0], Newtonsoft.Json.Formatting.Indented) : "null";
|
||||
Assert.Fail($"GetMessagesAsync returned 0 messages, but direct count query found {count} items for conversation {conversationId}. Raw document: {rawJson}");
|
||||
}
|
||||
|
||||
Assert.Single(messageList);
|
||||
Assert.Equal("Hello, world!", messageList[0].Text);
|
||||
Assert.Equal(ChatRole.User, messageList[0].Role);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var conversationId = Guid.NewGuid().ToString();
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId);
|
||||
var messages = new[]
|
||||
{
|
||||
new ChatMessage(ChatRole.User, "First message"),
|
||||
new ChatMessage(ChatRole.Assistant, "Second message"),
|
||||
new ChatMessage(ChatRole.User, "Third message")
|
||||
};
|
||||
|
||||
// Act
|
||||
await store.AddMessagesAsync(messages);
|
||||
|
||||
// Assert
|
||||
var retrievedMessages = await store.GetMessagesAsync();
|
||||
var messageList = retrievedMessages.ToList();
|
||||
Assert.Equal(3, messageList.Count);
|
||||
Assert.Equal("First message", messageList[0].Text);
|
||||
Assert.Equal("Second message", messageList[1].Text);
|
||||
Assert.Equal("Third message", messageList[2].Text);
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region GetMessagesAsync Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString());
|
||||
|
||||
// Act
|
||||
var messages = await store.GetMessagesAsync();
|
||||
|
||||
// Assert
|
||||
Assert.Empty(messages);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var conversation1 = Guid.NewGuid().ToString();
|
||||
var conversation2 = Guid.NewGuid().ToString();
|
||||
|
||||
using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation1);
|
||||
using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation2);
|
||||
|
||||
await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]);
|
||||
await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]);
|
||||
|
||||
// Act
|
||||
var messages1 = await store1.GetMessagesAsync();
|
||||
var messages2 = await store2.GetMessagesAsync();
|
||||
|
||||
// Assert
|
||||
var messageList1 = messages1.ToList();
|
||||
var messageList2 = messages2.ToList();
|
||||
Assert.Single(messageList1);
|
||||
Assert.Single(messageList2);
|
||||
Assert.Equal("Message for conversation 1", messageList1[0].Text);
|
||||
Assert.Equal("Message for conversation 2", messageList2[0].Text);
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Integration Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID
|
||||
using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId);
|
||||
|
||||
var messages = new[]
|
||||
{
|
||||
new ChatMessage(ChatRole.System, "You are a helpful assistant."),
|
||||
new ChatMessage(ChatRole.User, "Hello!"),
|
||||
new ChatMessage(ChatRole.Assistant, "Hi there! How can I help you today?"),
|
||||
new ChatMessage(ChatRole.User, "What's the weather like?"),
|
||||
new ChatMessage(ChatRole.Assistant, "I'm sorry, I don't have access to current weather data.")
|
||||
};
|
||||
|
||||
// Act 1: Add messages
|
||||
await originalStore.AddMessagesAsync(messages);
|
||||
|
||||
// Act 2: Verify messages were added
|
||||
var retrievedMessages = await originalStore.GetMessagesAsync();
|
||||
var retrievedList = retrievedMessages.ToList();
|
||||
Assert.Equal(5, retrievedList.Count);
|
||||
|
||||
// Act 3: Create new store instance for same conversation (test persistence)
|
||||
using var newStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId);
|
||||
var persistedMessages = await newStore.GetMessagesAsync();
|
||||
var persistedList = persistedMessages.ToList();
|
||||
|
||||
// Assert final state
|
||||
Assert.Equal(5, persistedList.Count);
|
||||
Assert.Equal("You are a helpful assistant.", persistedList[0].Text);
|
||||
Assert.Equal("Hello!", persistedList[1].Text);
|
||||
Assert.Equal("Hi there! How can I help you today?", persistedList[2].Text);
|
||||
Assert.Equal("What's the weather like?", persistedList[3].Text);
|
||||
Assert.Equal("I'm sorry, I don't have access to current weather data.", persistedList[4].Text);
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Disposal Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Dispose_AfterUse_ShouldNotThrow()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString());
|
||||
|
||||
// Act & Assert
|
||||
store.Dispose(); // Should not throw
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Dispose_MultipleCalls_ShouldNotThrow()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString());
|
||||
|
||||
// Act & Assert
|
||||
store.Dispose(); // First call
|
||||
store.Dispose(); // Second call - should not throw
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Hierarchical Partitioning Tests
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance()
|
||||
{
|
||||
// Arrange & Act
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789");
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(store);
|
||||
Assert.Equal("session-789", store.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(HierarchicalTestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance()
|
||||
{
|
||||
// Arrange & Act
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
TokenCredential credential = new DefaultAzureCredential();
|
||||
using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789");
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(store);
|
||||
Assert.Equal("session-789", store.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(HierarchicalTestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance()
|
||||
{
|
||||
// Arrange & Act
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey);
|
||||
using var store = new CosmosChatMessageStore(cosmosClient, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789");
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(store);
|
||||
Assert.Equal("session-789", store.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(HierarchicalTestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException()
|
||||
{
|
||||
// Arrange & Act & Assert
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
Assert.Throws<ArgumentNullException>(() =>
|
||||
new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, null!, "user-456", "session-789"));
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException()
|
||||
{
|
||||
// Arrange & Act & Assert
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
Assert.Throws<ArgumentException>(() =>
|
||||
new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789"));
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException()
|
||||
{
|
||||
// Arrange & Act & Assert
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
Assert.Throws<ArgumentException>(() =>
|
||||
new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " "));
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string TenantId = "tenant-123";
|
||||
const string UserId = "user-456";
|
||||
const string SessionId = "session-789";
|
||||
// Test hierarchical partitioning constructor with connection string
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId);
|
||||
var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!");
|
||||
|
||||
// Act
|
||||
await store.AddMessagesAsync([message]);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Assert
|
||||
var messages = await store.GetMessagesAsync();
|
||||
var messageList = messages.ToList();
|
||||
|
||||
Assert.Single(messageList);
|
||||
Assert.Equal("Hello from hierarchical partitioning!", messageList[0].Text);
|
||||
Assert.Equal(ChatRole.User, messageList[0].Role);
|
||||
|
||||
// Verify that the document is stored with hierarchical partitioning metadata
|
||||
var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type")
|
||||
.WithParameter("@conversationId", SessionId)
|
||||
.WithParameter("@type", "ChatMessage");
|
||||
|
||||
var iterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(HierarchicalTestContainerId)
|
||||
.GetItemQueryIterator<dynamic>(directQuery, requestOptions: new QueryRequestOptions
|
||||
{
|
||||
PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build()
|
||||
});
|
||||
|
||||
var response = await iterator.ReadNextAsync();
|
||||
var document = response.FirstOrDefault();
|
||||
|
||||
Assert.NotNull(document);
|
||||
// The document should have hierarchical metadata
|
||||
Assert.Equal(SessionId, (string)document!.conversationId);
|
||||
Assert.Equal(TenantId, (string)document!.tenantId);
|
||||
Assert.Equal(UserId, (string)document!.userId);
|
||||
Assert.Equal(SessionId, (string)document!.sessionId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string TenantId = "tenant-batch";
|
||||
const string UserId = "user-batch";
|
||||
const string SessionId = "session-batch";
|
||||
// Test hierarchical partitioning constructor with connection string
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId);
|
||||
var messages = new[]
|
||||
{
|
||||
new ChatMessage(ChatRole.User, "First hierarchical message"),
|
||||
new ChatMessage(ChatRole.Assistant, "Second hierarchical message"),
|
||||
new ChatMessage(ChatRole.User, "Third hierarchical message")
|
||||
};
|
||||
|
||||
// Act
|
||||
await store.AddMessagesAsync(messages);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Assert
|
||||
var retrievedMessages = await store.GetMessagesAsync();
|
||||
var messageList = retrievedMessages.ToList();
|
||||
|
||||
Assert.Equal(3, messageList.Count);
|
||||
Assert.Equal("First hierarchical message", messageList[0].Text);
|
||||
Assert.Equal("Second hierarchical message", messageList[1].Text);
|
||||
Assert.Equal("Third hierarchical message", messageList[2].Text);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string TenantId = "tenant-isolation";
|
||||
const string UserId1 = "user-1";
|
||||
const string UserId2 = "user-2";
|
||||
const string SessionId = "session-isolation";
|
||||
|
||||
// Different userIds create different hierarchical partitions, providing proper isolation
|
||||
using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId);
|
||||
using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId);
|
||||
|
||||
// Add messages to both stores
|
||||
await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]);
|
||||
await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 2")]);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Act & Assert
|
||||
var messages1 = await store1.GetMessagesAsync();
|
||||
var messageList1 = messages1.ToList();
|
||||
|
||||
var messages2 = await store2.GetMessagesAsync();
|
||||
var messageList2 = messages2.ToList();
|
||||
|
||||
// With true hierarchical partitioning, each user sees only their own messages
|
||||
Assert.Single(messageList1);
|
||||
Assert.Single(messageList2);
|
||||
Assert.Equal("Message from user 1", messageList1[0].Text);
|
||||
Assert.Equal("Message from user 2", messageList2[0].Text);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string TenantId = "tenant-serialize";
|
||||
const string UserId = "user-serialize";
|
||||
const string SessionId = "session-serialize";
|
||||
|
||||
using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId);
|
||||
await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]);
|
||||
|
||||
// Act - Serialize the store state
|
||||
var serializedState = originalStore.Serialize();
|
||||
|
||||
// Create a new store from the serialized state
|
||||
using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey);
|
||||
var serializerOptions = new JsonSerializerOptions
|
||||
{
|
||||
TypeInfoResolver = new DefaultJsonTypeInfoResolver()
|
||||
};
|
||||
using var deserializedStore = CosmosChatMessageStore.CreateFromSerializedState(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Assert - The deserialized store should have the same functionality
|
||||
var messages = await deserializedStore.GetMessagesAsync();
|
||||
var messageList = messages.ToList();
|
||||
|
||||
Assert.Single(messageList);
|
||||
Assert.Equal("Test serialization message", messageList[0].Text);
|
||||
Assert.Equal(SessionId, deserializedStore.ConversationId);
|
||||
Assert.Equal(s_testDatabaseId, deserializedStore.DatabaseId);
|
||||
Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string SessionId = "coexist-session";
|
||||
|
||||
// Create simple store using simple partitioning container and hierarchical store using hierarchical container
|
||||
using var simpleStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, SessionId);
|
||||
using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId);
|
||||
|
||||
// Add messages to both
|
||||
await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]);
|
||||
await hierarchicalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]);
|
||||
|
||||
// Wait a moment for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Act & Assert
|
||||
var simpleMessages = await simpleStore.GetMessagesAsync();
|
||||
var simpleMessageList = simpleMessages.ToList();
|
||||
|
||||
var hierarchicalMessages = await hierarchicalStore.GetMessagesAsync();
|
||||
var hierarchicalMessageList = hierarchicalMessages.ToList();
|
||||
|
||||
// Each should only see its own messages since they use different containers
|
||||
Assert.Single(simpleMessageList);
|
||||
Assert.Single(hierarchicalMessageList);
|
||||
Assert.Equal("Simple partitioning message", simpleMessageList[0].Text);
|
||||
Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string ConversationId = "max-messages-test";
|
||||
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId);
|
||||
|
||||
// Add 10 messages
|
||||
var messages = new List<ChatMessage>();
|
||||
for (int i = 1; i <= 10; i++)
|
||||
{
|
||||
messages.Add(new ChatMessage(ChatRole.User, $"Message {i}"));
|
||||
await Task.Delay(10); // Small delay to ensure different timestamps
|
||||
}
|
||||
await store.AddMessagesAsync(messages);
|
||||
|
||||
// Wait for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Act - Set max to 5 and retrieve
|
||||
store.MaxMessagesToRetrieve = 5;
|
||||
var retrievedMessages = await store.GetMessagesAsync();
|
||||
var messageList = retrievedMessages.ToList();
|
||||
|
||||
// Assert - Should get the 5 most recent messages (6-10) in ascending order
|
||||
Assert.Equal(5, messageList.Count);
|
||||
Assert.Equal("Message 6", messageList[0].Text);
|
||||
Assert.Equal("Message 7", messageList[1].Text);
|
||||
Assert.Equal("Message 8", messageList[2].Text);
|
||||
Assert.Equal("Message 9", messageList[3].Text);
|
||||
Assert.Equal("Message 10", messageList[4].Text);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
[Trait("Category", "CosmosDB")]
|
||||
public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
const string ConversationId = "max-messages-null-test";
|
||||
|
||||
using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId);
|
||||
|
||||
// Add 10 messages
|
||||
var messages = new List<ChatMessage>();
|
||||
for (int i = 1; i <= 10; i++)
|
||||
{
|
||||
messages.Add(new ChatMessage(ChatRole.User, $"Message {i}"));
|
||||
}
|
||||
await store.AddMessagesAsync(messages);
|
||||
|
||||
// Wait for eventual consistency
|
||||
await Task.Delay(100);
|
||||
|
||||
// Act - No limit set (default null)
|
||||
var retrievedMessages = await store.GetMessagesAsync();
|
||||
var messageList = retrievedMessages.ToList();
|
||||
|
||||
// Assert - Should get all 10 messages
|
||||
Assert.Equal(10, messageList.Count);
|
||||
Assert.Equal("Message 1", messageList[0].Text);
|
||||
Assert.Equal("Message 10", messageList[9].Text);
|
||||
}
|
||||
|
||||
#endregion
|
||||
}
|
||||
@@ -0,0 +1,454 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System;
|
||||
using System.Linq;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Tasks;
|
||||
using Microsoft.Agents.AI.Workflows;
|
||||
using Microsoft.Agents.AI.Workflows.Checkpointing;
|
||||
using Microsoft.Azure.Cosmos;
|
||||
using Xunit;
|
||||
|
||||
namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests;
|
||||
|
||||
/// <summary>
|
||||
/// Contains tests for <see cref="CosmosCheckpointStore"/>.
|
||||
///
|
||||
/// Test Modes:
|
||||
/// - Default Mode: Cleans up all test data after each test run (deletes database)
|
||||
/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer
|
||||
///
|
||||
/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true
|
||||
/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test
|
||||
///
|
||||
/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at:
|
||||
/// https://localhost:8081/_explorer/index.html
|
||||
/// Database: AgentFrameworkTests
|
||||
/// Container: Checkpoints
|
||||
/// </summary>
|
||||
[Collection("CosmosDB")]
|
||||
public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable
|
||||
{
|
||||
// Cosmos DB Emulator connection settings
|
||||
private const string EmulatorEndpoint = "https://localhost:8081";
|
||||
private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==";
|
||||
private const string TestContainerId = "Checkpoints";
|
||||
// Use unique database ID per test class instance to avoid conflicts
|
||||
#pragma warning disable CA1802 // Use literals where appropriate
|
||||
private static readonly string s_testDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}";
|
||||
#pragma warning restore CA1802
|
||||
|
||||
private string _connectionString = string.Empty;
|
||||
private CosmosClient? _cosmosClient;
|
||||
private Database? _database;
|
||||
private bool _emulatorAvailable;
|
||||
private bool _preserveContainer;
|
||||
|
||||
// JsonSerializerOptions configured for .NET 9+ compatibility
|
||||
private static readonly JsonSerializerOptions s_jsonOptions = CreateJsonOptions();
|
||||
|
||||
private static JsonSerializerOptions CreateJsonOptions()
|
||||
{
|
||||
var options = new JsonSerializerOptions();
|
||||
#if NET9_0_OR_GREATER
|
||||
options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver();
|
||||
#endif
|
||||
return options;
|
||||
}
|
||||
|
||||
public async Task InitializeAsync()
|
||||
{
|
||||
// Check environment variable to determine if we should preserve containers
|
||||
// Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection
|
||||
this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}";
|
||||
|
||||
try
|
||||
{
|
||||
this._cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey);
|
||||
|
||||
// Test connection by attempting to create database
|
||||
this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId);
|
||||
await this._database.CreateContainerIfNotExistsAsync(
|
||||
TestContainerId,
|
||||
"/runId",
|
||||
throughput: 400);
|
||||
|
||||
this._emulatorAvailable = true;
|
||||
}
|
||||
catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException))
|
||||
{
|
||||
// Emulator not available, tests will be skipped
|
||||
this._emulatorAvailable = false;
|
||||
this._cosmosClient?.Dispose();
|
||||
this._cosmosClient = null;
|
||||
}
|
||||
}
|
||||
|
||||
public async Task DisposeAsync()
|
||||
{
|
||||
if (this._cosmosClient != null && this._emulatorAvailable)
|
||||
{
|
||||
try
|
||||
{
|
||||
if (this._preserveContainer)
|
||||
{
|
||||
// Preserve mode: Don't delete the database/container, keep data for inspection
|
||||
// This allows viewing data in the Cosmos DB Emulator Data Explorer
|
||||
// No cleanup needed - data persists for debugging
|
||||
}
|
||||
else
|
||||
{
|
||||
// Clean mode: Delete the test database and all data
|
||||
await this._database!.DeleteAsync();
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
// Ignore cleanup errors, but log for diagnostics
|
||||
Console.WriteLine($"[DisposeAsync] Cleanup error: {ex.Message}\n{ex.StackTrace}");
|
||||
}
|
||||
finally
|
||||
{
|
||||
this._cosmosClient.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void SkipIfEmulatorNotAvailable()
|
||||
{
|
||||
// In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true"
|
||||
// Locally: Skip if emulator connection check failed
|
||||
var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available");
|
||||
}
|
||||
|
||||
#region Constructor Tests
|
||||
|
||||
[SkippableFact]
|
||||
public void Constructor_WithCosmosClient_SetsProperties()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
|
||||
// Assert
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(TestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public void Constructor_WithConnectionString_SetsProperties()
|
||||
{
|
||||
// Arrange
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Act
|
||||
using var store = new CosmosCheckpointStore(this._connectionString, s_testDatabaseId, TestContainerId);
|
||||
|
||||
// Assert
|
||||
Assert.Equal(s_testDatabaseId, store.DatabaseId);
|
||||
Assert.Equal(TestContainerId, store.ContainerId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException()
|
||||
{
|
||||
// Act & Assert
|
||||
Assert.Throws<ArgumentNullException>(() =>
|
||||
new CosmosCheckpointStore((CosmosClient)null!, s_testDatabaseId, TestContainerId));
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public void Constructor_WithNullConnectionString_ThrowsArgumentException()
|
||||
{
|
||||
// Act & Assert
|
||||
Assert.Throws<ArgumentNullException>(() =>
|
||||
new CosmosCheckpointStore((string)null!, s_testDatabaseId, TestContainerId));
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Checkpoint Operations Tests
|
||||
|
||||
[SkippableFact]
|
||||
public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfullyAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions);
|
||||
|
||||
// Act
|
||||
var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(checkpointInfo);
|
||||
Assert.Equal(runId, checkpointInfo.RunId);
|
||||
Assert.NotNull(checkpointInfo.CheckpointId);
|
||||
Assert.NotEmpty(checkpointInfo.CheckpointId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValueAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow };
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions);
|
||||
|
||||
// Act
|
||||
var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
var retrievedValue = await store.RetrieveCheckpointAsync(runId, checkpointInfo);
|
||||
|
||||
// Assert
|
||||
Assert.Equal(JsonValueKind.Object, retrievedValue.ValueKind);
|
||||
Assert.True(retrievedValue.TryGetProperty("message", out var messageProp));
|
||||
Assert.Equal("Hello, World!", messageProp.GetString());
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationExceptionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint");
|
||||
|
||||
// Act & Assert
|
||||
await Assert.ThrowsAsync<InvalidOperationException>(() =>
|
||||
store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask());
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollectionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
|
||||
// Act
|
||||
var index = await store.RetrieveIndexAsync(runId);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(index);
|
||||
Assert.Empty(index);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpointsAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Create multiple checkpoints
|
||||
var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
var checkpoint2 = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
var checkpoint3 = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
|
||||
// Act
|
||||
var index = (await store.RetrieveIndexAsync(runId)).ToList();
|
||||
|
||||
// Assert
|
||||
Assert.Equal(3, index.Count);
|
||||
Assert.Contains(index, c => c.CheckpointId == checkpoint1.CheckpointId);
|
||||
Assert.Contains(index, c => c.CheckpointId == checkpoint2.CheckpointId);
|
||||
Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task CreateCheckpointAsync_WithParent_CreatesHierarchyAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Act
|
||||
var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
var childCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue, parentCheckpoint);
|
||||
|
||||
// Assert
|
||||
Assert.NotEqual(parentCheckpoint.CheckpointId, childCheckpoint.CheckpointId);
|
||||
Assert.Equal(runId, parentCheckpoint.RunId);
|
||||
Assert.Equal(runId, childCheckpoint.RunId);
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResultsAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Create parent and child checkpoints
|
||||
var parent = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
var child1 = await store.CreateCheckpointAsync(runId, checkpointValue, parent);
|
||||
var child2 = await store.CreateCheckpointAsync(runId, checkpointValue, parent);
|
||||
|
||||
// Create an orphan checkpoint
|
||||
var orphan = await store.CreateCheckpointAsync(runId, checkpointValue);
|
||||
|
||||
// Act
|
||||
var allCheckpoints = (await store.RetrieveIndexAsync(runId)).ToList();
|
||||
var childrenOfParent = (await store.RetrieveIndexAsync(runId, parent)).ToList();
|
||||
|
||||
// Assert
|
||||
Assert.Equal(4, allCheckpoints.Count); // parent + 2 children + orphan
|
||||
Assert.Equal(2, childrenOfParent.Count); // only children
|
||||
|
||||
Assert.Contains(childrenOfParent, c => c.CheckpointId == child1.CheckpointId);
|
||||
Assert.Contains(childrenOfParent, c => c.CheckpointId == child2.CheckpointId);
|
||||
Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == parent.CheckpointId);
|
||||
Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == orphan.CheckpointId);
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Run Isolation Tests
|
||||
|
||||
[SkippableFact]
|
||||
public async Task CheckpointOperations_DifferentRuns_IsolatesDataAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId1 = Guid.NewGuid().ToString();
|
||||
var runId2 = Guid.NewGuid().ToString();
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Act
|
||||
var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue);
|
||||
var checkpoint2 = await store.CreateCheckpointAsync(runId2, checkpointValue);
|
||||
|
||||
var index1 = (await store.RetrieveIndexAsync(runId1)).ToList();
|
||||
var index2 = (await store.RetrieveIndexAsync(runId2)).ToList();
|
||||
|
||||
// Assert
|
||||
Assert.Single(index1);
|
||||
Assert.Single(index2);
|
||||
Assert.Equal(checkpoint1.CheckpointId, index1[0].CheckpointId);
|
||||
Assert.Equal(checkpoint2.CheckpointId, index2[0].CheckpointId);
|
||||
Assert.NotEqual(checkpoint1.CheckpointId, checkpoint2.CheckpointId);
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Error Handling Tests
|
||||
|
||||
[SkippableFact]
|
||||
public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentExceptionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Act & Assert
|
||||
await Assert.ThrowsAsync<ArgumentException>(() =>
|
||||
store.CreateCheckpointAsync(null!, checkpointValue).AsTask());
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentExceptionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Act & Assert
|
||||
await Assert.ThrowsAsync<ArgumentException>(() =>
|
||||
store.CreateCheckpointAsync("", checkpointValue).AsTask());
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullExceptionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var runId = Guid.NewGuid().ToString();
|
||||
|
||||
// Act & Assert
|
||||
await Assert.ThrowsAsync<ArgumentNullException>(() =>
|
||||
store.RetrieveCheckpointAsync(runId, null!).AsTask());
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
#region Disposal Tests
|
||||
|
||||
[SkippableFact]
|
||||
public async Task Dispose_AfterDisposal_ThrowsObjectDisposedExceptionAsync()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions);
|
||||
|
||||
// Act
|
||||
store.Dispose();
|
||||
|
||||
// Assert
|
||||
await Assert.ThrowsAsync<ObjectDisposedException>(() =>
|
||||
store.CreateCheckpointAsync("test-run", checkpointValue).AsTask());
|
||||
}
|
||||
|
||||
[SkippableFact]
|
||||
public void Dispose_MultipleCalls_DoesNotThrow()
|
||||
{
|
||||
this.SkipIfEmulatorNotAvailable();
|
||||
|
||||
// Arrange
|
||||
var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId);
|
||||
|
||||
// Act & Assert (should not throw)
|
||||
store.Dispose();
|
||||
store.Dispose();
|
||||
store.Dispose();
|
||||
}
|
||||
|
||||
#endregion
|
||||
|
||||
public void Dispose()
|
||||
{
|
||||
this.Dispose(true);
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
|
||||
protected virtual void Dispose(bool disposing)
|
||||
{
|
||||
if (disposing)
|
||||
{
|
||||
this._cosmosClient?.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using Xunit;
|
||||
|
||||
namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests;
|
||||
|
||||
/// <summary>
|
||||
/// Defines a collection fixture for Cosmos DB tests to ensure they run sequentially.
|
||||
/// This prevents race conditions and resource conflicts when tests create and delete
|
||||
/// databases in the Cosmos DB Emulator.
|
||||
/// </summary>
|
||||
[CollectionDefinition("CosmosDB", DisableParallelization = true)]
|
||||
public sealed class CosmosDBCollectionFixture
|
||||
{
|
||||
// This class has no code, and is never created. Its purpose is simply
|
||||
// to be the place to apply [CollectionDefinition] and all the
|
||||
// ICollectionFixture<> interfaces.
|
||||
}
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
<Project Sdk="Microsoft.NET.Sdk">
|
||||
|
||||
<PropertyGroup>
|
||||
<TargetFrameworks>net10.0;net9.0</TargetFrameworks>
|
||||
<NoWarn>$(NoWarn);MEAI001</NoWarn>
|
||||
</PropertyGroup>
|
||||
|
||||
<PropertyGroup Condition="$([MSBuild]::IsTargetFrameworkCompatible('$(TargetFramework)', 'net8.0'))">
|
||||
<JsonSerializerIsReflectionEnabledByDefault>false</JsonSerializerIsReflectionEnabledByDefault>
|
||||
</PropertyGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<ProjectReference Include="..\..\src\Microsoft.Agents.AI.CosmosNoSql\Microsoft.Agents.AI.CosmosNoSql.csproj" />
|
||||
</ItemGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="System.Text.Json" />
|
||||
<PackageReference Include="System.Linq.AsyncEnumerable" />
|
||||
<PackageReference Include="Azure.Identity" />
|
||||
<PackageReference Include="Microsoft.Azure.Cosmos" />
|
||||
<PackageReference Include="Xunit.SkippableFact" />
|
||||
</ItemGroup>
|
||||
|
||||
</Project>
|
||||
Reference in New Issue
Block a user