// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; using System.Linq; using System.Text.Json; using System.Text.Json.Nodes; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.VectorData; using Microsoft.SemanticKernel.Connectors.CosmosNoSql; using Moq; using Xunit; using DistanceFunction = Microsoft.Extensions.VectorData.DistanceFunction; using IndexKind = Microsoft.Extensions.VectorData.IndexKind; namespace SemanticKernel.Connectors.CosmosNoSql.UnitTests; /// /// Unit tests for class. /// public sealed class CosmosNoSqlCollectionTests { private readonly Mock _mockDatabase = new(); private readonly Mock _mockContainer = new(); public CosmosNoSqlCollectionTests() { this._mockDatabase.Setup(l => l.GetContainer(It.IsAny())).Returns(this._mockContainer.Object); var mockClient = new Mock(); mockClient.Setup(l => l.ClientOptions).Returns(new CosmosClientOptions() { UseSystemTextJsonSerializerWithOptions = JsonSerializerOptions.Default }); this._mockDatabase .Setup(l => l.Client) .Returns(mockClient.Object); } [Fact] public void ConstructorForModelWithoutKeyThrowsException() { // Act & Assert var exception = Assert.Throws(() => new CosmosNoSqlCollection(this._mockDatabase.Object, "collection")); Assert.Contains("No key property found", exception.Message); } [Fact] public void ConstructorWithoutSystemTextJsonSerializerOptionsThrowsArgumentException() { // Arrange var mockDatabase = new Mock(); var mockClient = new Mock(); mockDatabase.Setup(l => l.Client).Returns(mockClient.Object); // Act & Assert var exception = Assert.Throws(() => new CosmosNoSqlCollection(mockDatabase.Object, "collection")); Assert.Contains(nameof(CosmosClientOptions.UseSystemTextJsonSerializerWithOptions), exception.Message); } [Fact] public void ConstructorWithDeclarativeModelInitializesCollection() { // Act & Assert using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); Assert.NotNull(collection); } [Fact] public void ConstructorWithImperativeModelInitializesCollection() { // Arrange var definition = new VectorStoreCollectionDefinition { Properties = [new VectorStoreKeyProperty("Id", typeof(string))] }; // Act using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", new() { Definition = definition }); // Assert Assert.NotNull(collection); } [Theory] [MemberData(nameof(CollectionExistsData))] public async Task CollectionExistsReturnsValidResultAsync(List collections, string collectionName, bool expectedResult) { // Arrange var mockFeedResponse = new Mock>(); mockFeedResponse .Setup(l => l.Resource) .Returns(collections); var mockFeedIterator = new Mock>(); mockFeedIterator .SetupSequence(l => l.HasMoreResults) .Returns(true) .Returns(false); mockFeedIterator .Setup(l => l.ReadNextAsync(It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); this._mockDatabase .Setup(l => l.GetContainerQueryIterator( It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockFeedIterator.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, collectionName); // Act var actualResult = await sut.CollectionExistsAsync(); // Assert Assert.Equal(expectedResult, actualResult); } [Theory] [InlineData(IndexingMode.Consistent)] [InlineData(IndexingMode.Lazy)] [InlineData(IndexingMode.None)] public async Task EnsureCollectionExistsUsesValidContainerPropertiesAsync(IndexingMode indexingMode) { // Arrange const string CollectionName = "collection"; var mockFeedResponse = new Mock>(); mockFeedResponse .Setup(l => l.Resource) .Returns([]); var mockFeedIterator = new Mock>(); mockFeedIterator .SetupSequence(l => l.HasMoreResults) .Returns(true) .Returns(false); mockFeedIterator .Setup(l => l.ReadNextAsync(It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); this._mockDatabase .Setup(l => l.GetContainerQueryIterator( It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockFeedIterator.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, CollectionName, new() { IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None }); var expectedVectorEmbeddingPolicy = new VectorEmbeddingPolicy( [ new Embedding { DataType = VectorDataType.Float32, Dimensions = 2, DistanceFunction = Microsoft.Azure.Cosmos.DistanceFunction.Cosine, Path = "/DescriptionEmbedding2" }, new Embedding { DataType = VectorDataType.Uint8, Dimensions = 3, DistanceFunction = Microsoft.Azure.Cosmos.DistanceFunction.DotProduct, Path = "/DescriptionEmbedding3" }, new Embedding { DataType = VectorDataType.Int8, Dimensions = 4, DistanceFunction = Microsoft.Azure.Cosmos.DistanceFunction.Euclidean, Path = "/DescriptionEmbedding4" }, ]); var expectedIndexingPolicy = new IndexingPolicy { VectorIndexes = [ new VectorIndexPath { Type = VectorIndexType.Flat, Path = "/DescriptionEmbedding2" }, new VectorIndexPath { Type = VectorIndexType.QuantizedFlat, Path = "/DescriptionEmbedding3" }, new VectorIndexPath { Type = VectorIndexType.DiskANN, Path = "/DescriptionEmbedding4" }, ], IndexingMode = indexingMode, Automatic = indexingMode != IndexingMode.None }; if (indexingMode != IndexingMode.None) { expectedIndexingPolicy.IncludedPaths.Add(new IncludedPath { Path = "/IndexableData1/?" }); expectedIndexingPolicy.IncludedPaths.Add(new IncludedPath { Path = "/IndexableData2/?" }); expectedIndexingPolicy.IncludedPaths.Add(new IncludedPath { Path = "/" }); expectedIndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/DescriptionEmbedding2/*" }); expectedIndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/DescriptionEmbedding3/*" }); expectedIndexingPolicy.ExcludedPaths.Add(new ExcludedPath { Path = "/DescriptionEmbedding4/*" }); } var expectedContainerProperties = new ContainerProperties(CollectionName, "/id") { VectorEmbeddingPolicy = expectedVectorEmbeddingPolicy, IndexingPolicy = expectedIndexingPolicy }; // Act await sut.EnsureCollectionExistsAsync(); // Assert this._mockDatabase.Verify(l => l.CreateContainerAsync( It.Is(properties => this.VerifyContainerProperties(expectedContainerProperties, properties)), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once()); } [Theory] [MemberData(nameof(EnsureCollectionExistsData))] public async Task EnsureCollectionExistsInvokesValidMethodsAsync(List collections, int actualCollectionCreations) { // Arrange const string CollectionName = "collection"; var mockFeedResponse = new Mock>(); mockFeedResponse .Setup(l => l.Resource) .Returns(collections); var mockFeedIterator = new Mock>(); mockFeedIterator .SetupSequence(l => l.HasMoreResults) .Returns(true) .Returns(false); mockFeedIterator .Setup(l => l.ReadNextAsync(It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); this._mockDatabase .Setup(l => l.GetContainerQueryIterator( It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockFeedIterator.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, CollectionName); // Act await sut.EnsureCollectionExistsAsync(); // Assert this._mockDatabase.Verify(l => l.CreateContainerAsync( It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny()), Times.Exactly(actualCollectionCreations)); } [Theory] [InlineData("recordKey", false)] [InlineData("partitionKey", true)] public async Task DeleteInvokesValidMethodsAsync( string expectedPartitionKey, bool useExplicitPartitionKey) { // Arrange const string RecordKey = "recordKey"; const string PartitionKey = "partitionKey"; // Act if (useExplicitPartitionKey) { using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); await ((VectorStoreCollection)sut).DeleteAsync( new CosmosNoSqlKey(RecordKey, PartitionKey)); } else { using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); await ((VectorStoreCollection)sut).DeleteAsync( RecordKey); } // Assert this._mockContainer.Verify(l => l.DeleteItemAsync( RecordKey, new PartitionKey(expectedPartitionKey), It.IsAny(), It.IsAny()), Times.Once()); } [Fact] public async Task DeleteBatchInvokesValidMethodsAsync() { // Arrange List recordKeys = ["key1", "key2"]; using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act await sut.DeleteAsync(recordKeys); // Assert foreach (var key in recordKeys) { this._mockContainer.Verify(l => l.DeleteItemAsync( key, new PartitionKey(key), It.IsAny(), It.IsAny()), Times.Once()); } } [Fact] public async Task DeleteCollectionInvokesValidMethodsAsync() { // Arrange using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act await sut.EnsureCollectionDeletedAsync(); // Assert this._mockContainer.Verify(l => l.DeleteContainerAsync( It.IsAny(), It.IsAny()), Times.Once()); } [Fact] public async Task GetReturnsValidRecordAsync() { // Arrange const string RecordKey = "key"; var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name" }; var mockItemResponse = new Mock>(); mockItemResponse .Setup(l => l.Resource) .Returns(jsonObject); this._mockContainer .Setup(l => l.ReadItemAsync( RecordKey, new PartitionKey(RecordKey), It.IsAny(), It.IsAny())) .ReturnsAsync(mockItemResponse.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act var result = await sut.GetAsync(RecordKey); // Assert Assert.NotNull(result); Assert.Equal(RecordKey, result.HotelId); Assert.Equal("Test Name", result.HotelName); } [Fact] public async Task GetBatchReturnsValidRecordAsync() { // Arrange var jsonObject1 = new JsonObject { ["id"] = "key1", ["HotelName"] = "Test Name 1" }; var jsonObject2 = new JsonObject { ["id"] = "key2", ["HotelName"] = "Test Name 2" }; var jsonObject3 = new JsonObject { ["id"] = "key3", ["HotelName"] = "Test Name 3" }; var mockFeedResponse = new Mock>(); mockFeedResponse .Setup(l => l.Resource) .Returns([jsonObject1, jsonObject2, jsonObject3]); this._mockContainer .Setup(l => l.ReadManyItemsAsync( It.IsAny>(), It.IsAny(), It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act var results = await sut.GetAsync(["key1", "key2", "key3"]).ToListAsync(); // Assert Assert.NotNull(results[0]); Assert.Equal("key1", results[0].HotelId); Assert.Equal("Test Name 1", results[0].HotelName); Assert.NotNull(results[1]); Assert.Equal("key2", results[1].HotelId); Assert.Equal("Test Name 2", results[1].HotelName); Assert.NotNull(results[2]); Assert.Equal("key3", results[2].HotelId); Assert.Equal("Test Name 3", results[2].HotelName); } [Fact] public async Task CanUpsertRecordAsync() { // Arrange var hotel = new CosmosNoSqlHotel("key") { HotelName = "Test Name" }; using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act await sut.UpsertAsync(hotel); // Assert this._mockContainer.Verify(l => l.UpsertItemAsync( It.Is(node => node["id"]!.ToString() == "key" && node["HotelName"]!.ToString() == "Test Name"), new PartitionKey("key"), It.IsAny(), It.IsAny()), Times.Once()); } [Fact] public async Task CanUpsertManyRecordsAsync() { // Arrange var hotel1 = new CosmosNoSqlHotel("key1") { HotelName = "Test Name 1" }; var hotel2 = new CosmosNoSqlHotel("key2") { HotelName = "Test Name 2" }; var hotel3 = new CosmosNoSqlHotel("key3") { HotelName = "Test Name 3" }; using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act await sut.UpsertAsync([hotel1, hotel2, hotel3]); } [Fact] public async Task VectorizedSearchReturnsValidRecordAsync() { // Arrange const string RecordKey = "key"; const double ExpectedScore = 0.99; var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name", ["SimilarityScore"] = ExpectedScore }; var mockFeedResponse = new Mock>(); mockFeedResponse .Setup(l => l.Resource) .Returns([jsonObject]); var mockFeedIterator = new Mock>(); mockFeedIterator .SetupSequence(l => l.HasMoreResults) .Returns(true) .Returns(false); mockFeedIterator .Setup(l => l.ReadNextAsync(It.IsAny())) .ReturnsAsync(mockFeedResponse.Object); this._mockContainer .Setup(l => l.GetItemQueryIterator( It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockFeedIterator.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act var results = await sut.SearchAsync(new ReadOnlyMemory([1f, 2f, 3f]), top: 3).ToListAsync(); var result = results[0]; // Assert Assert.NotNull(result); Assert.Equal(RecordKey, result.Record.HotelId); Assert.Equal("Test Name", result.Record.HotelName); Assert.Equal(ExpectedScore, result.Score); } [Fact] public async Task VectorizedSearchWithUnsupportedVectorTypeThrowsExceptionAsync() { // Arrange using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act & Assert await Assert.ThrowsAsync(async () => await sut.SearchAsync(new List([1, 2, 3]), top: 3).ToListAsync()); } [Fact] public async Task VectorizedSearchWithNonExistentVectorPropertyNameThrowsExceptionAsync() { // Arrange using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); var searchOptions = new VectorSearchOptions { VectorProperty = r => "non-existent-property" }; // Act & Assert await Assert.ThrowsAsync(async () => await sut.SearchAsync(new ReadOnlyMemory([1f, 2f, 3f]), top: 3, searchOptions).ToListAsync()); } public static TheoryData, string, bool> CollectionExistsData => new() { { ["collection-2"], "collection-2", true }, { [], "non-existent-collection", false } }; public static TheoryData, int> EnsureCollectionExistsData => new() { { ["collection"], 0 }, { [], 1 } }; #region private bool VerifyContainerProperties(ContainerProperties expected, ContainerProperties actual) { Assert.Equal(expected.Id, actual.Id); Assert.Equal(expected.PartitionKeyPath, actual.PartitionKeyPath); Assert.Equal(expected.IndexingPolicy.IndexingMode, actual.IndexingPolicy.IndexingMode); Assert.Equal(expected.IndexingPolicy.Automatic, actual.IndexingPolicy.Automatic); if (expected.IndexingPolicy.IndexingMode != IndexingMode.None) { for (var i = 0; i < expected.VectorEmbeddingPolicy.Embeddings.Count; i++) { var expectedEmbedding = expected.VectorEmbeddingPolicy.Embeddings[i]; var actualEmbedding = actual.VectorEmbeddingPolicy.Embeddings[i]; Assert.Equal(expectedEmbedding.DataType, actualEmbedding.DataType); Assert.Equal(expectedEmbedding.Dimensions, actualEmbedding.Dimensions); Assert.Equal(expectedEmbedding.DistanceFunction, actualEmbedding.DistanceFunction); Assert.Equal(expectedEmbedding.Path, actualEmbedding.Path); } for (var i = 0; i < expected.IndexingPolicy.VectorIndexes.Count; i++) { var expectedIndexPath = expected.IndexingPolicy.VectorIndexes[i]; var actualIndexPath = actual.IndexingPolicy.VectorIndexes[i]; Assert.Equal(expectedIndexPath.Type, actualIndexPath.Type); Assert.Equal(expectedIndexPath.Path, actualIndexPath.Path); } for (var i = 0; i < expected.IndexingPolicy.IncludedPaths.Count; i++) { var expectedIncludedPath = expected.IndexingPolicy.IncludedPaths[i].Path; var actualIncludedPath = actual.IndexingPolicy.IncludedPaths[i].Path; Assert.Equal(expectedIncludedPath, actualIncludedPath); } for (var i = 0; i < expected.IndexingPolicy.ExcludedPaths.Count; i++) { var expectedExcludedPath = expected.IndexingPolicy.ExcludedPaths[i].Path; var actualExcludedPath = actual.IndexingPolicy.ExcludedPaths[i].Path; Assert.Equal(expectedExcludedPath, actualExcludedPath); } } return true; } #pragma warning disable CA1812 private sealed class TestModel { public string? Id { get; set; } public string? HotelName { get; set; } } private sealed class TestIndexingModel { [VectorStoreKey] public string? Id { get; set; } [VectorStoreVector(Dimensions: 2, DistanceFunction = DistanceFunction.CosineSimilarity, IndexKind = IndexKind.Flat)] public ReadOnlyMemory? DescriptionEmbedding2 { get; set; } [VectorStoreVector(Dimensions: 3, DistanceFunction = DistanceFunction.DotProductSimilarity, IndexKind = IndexKind.QuantizedFlat)] public ReadOnlyMemory? DescriptionEmbedding3 { get; set; } [VectorStoreVector(Dimensions: 4, DistanceFunction = DistanceFunction.EuclideanDistance, IndexKind = IndexKind.DiskAnn)] public ReadOnlyMemory? DescriptionEmbedding4 { get; set; } [VectorStoreData(IsIndexed = true)] public string? IndexableData1 { get; set; } [VectorStoreData(IsFullTextIndexed = true)] public string? IndexableData2 { get; set; } [VectorStoreData] public string? NonIndexableData1 { get; set; } } #pragma warning restore CA1812 [Fact] public void ConstructorWithStringKeyInitializesCollection() { // String TKey should work; partition key auto-defaults to the key property. using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); Assert.NotNull(collection); } [Fact] public void ConstructorWithStringKeyAndExplicitPartitionKeyInitializesCollection() { // String TKey should work when partition key explicitly points to the key property. using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", new() { PartitionKeyProperties = ["HotelId"] }); Assert.NotNull(collection); } [Fact] public void ConstructorWithStringKeyRejectsNonKeyPartitionKey() { // String TKey requires the partition key to be the key property. var exception = Assert.Throws(() => new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", new() { PartitionKeyProperties = ["HotelName"] })); Assert.Contains("partition key must be the key property", exception.Message); } [Fact] public void ConstructorWithStringKeyRejectsEmptyPartitionKey() { // String TKey cannot use empty partition key (PartitionKey.None) - that requires CosmosNoSqlKey. var exception = Assert.Throws(() => new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", new() { PartitionKeyProperties = [] })); Assert.Contains("partition key must be the key property", exception.Message); } [Fact] public void ConstructorWithGuidKeyInitializesCollection() { // Guid TKey should work; partition key auto-defaults to the key property. using var collection = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); Assert.NotNull(collection); } [Fact] public void ConstructorWithGuidKeyRejectsNonKeyPartitionKey() { // Guid TKey requires the partition key to be the key property. var exception = Assert.Throws(() => new CosmosNoSqlCollection( this._mockDatabase.Object, "collection", new() { PartitionKeyProperties = ["Name"] })); Assert.Contains("partition key must be the key property", exception.Message); } [Fact] public void ConstructorWithUnsupportedKeyTypeThrowsException() { var exception = Assert.Throws(() => new CosmosNoSqlCollection( this._mockDatabase.Object, "collection")); Assert.Contains("string, Guid, or CosmosNoSqlKey", exception.Message); } [Fact] public async Task DeleteWithStringKeyInvokesValidMethodsAsync() { // Arrange const string RecordKey = "recordKey"; using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act await sut.DeleteAsync(RecordKey); // Assert: with string key, partition key = document id this._mockContainer.Verify(l => l.DeleteItemAsync( RecordKey, new PartitionKey(RecordKey), It.IsAny(), It.IsAny()), Times.Once()); } [Fact] public async Task GetWithStringKeyReturnsValidRecordAsync() { // Arrange const string RecordKey = "key"; var jsonObject = new JsonObject { ["id"] = RecordKey, ["HotelName"] = "Test Name" }; var mockItemResponse = new Mock>(); mockItemResponse .Setup(l => l.Resource) .Returns(jsonObject); this._mockContainer .Setup(l => l.ReadItemAsync( RecordKey, new PartitionKey(RecordKey), It.IsAny(), It.IsAny())) .ReturnsAsync(mockItemResponse.Object); using var sut = new CosmosNoSqlCollection( this._mockDatabase.Object, "collection"); // Act var result = await sut.GetAsync(RecordKey); // Assert Assert.NotNull(result); Assert.Equal(RecordKey, result.HotelId); Assert.Equal("Test Name", result.HotelName); } #pragma warning disable CA1812 private sealed class GuidKeyModel { [VectorStoreKey] public Guid Id { get; set; } [VectorStoreData] public string? Name { get; set; } } #pragma warning restore CA1812 #endregion }