.Net: Azure CosmosDB MongoDB vCore Memory Store Bug Fixes

dotnet/src/Connectors/Connectors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStore.cs

@@ -22,6 +22,7 @@ public class AzureCosmosDBMongoDBMemoryStore : IMemoryStore, IDisposable
     private readonly MongoClient _mongoClient;
     private readonly IMongoDatabase _mongoDatabase;
     private readonly AzureCosmosDBMongoDBConfig _config;
+    private readonly bool _ownsMongoClient;
 
     /// <summary>
     /// Initiates a AzureCosmosDBMongoDBMemoryStore instance using a Azure CosmosDB Mongo vCore connection string
@@ -41,22 +42,21 @@ AzureCosmosDBMongoDBConfig config
         settings.ApplicationName = this._config.ApplicationName;
         this._mongoClient = new MongoClient(settings);
         this._mongoDatabase = this._mongoClient.GetDatabase(databaseName);
+        this._ownsMongoClient = true;
     }
 
     /// <summary>
     /// Initiates a AzureCosmosDBMongoDBMemoryStore instance using a Azure CosmosDB MongoDB client
     /// and other properties required for vector search.
     /// </summary>
     public AzureCosmosDBMongoDBMemoryStore(
-        IMongoClient mongoClient,
+        MongoClient mongoClient,
         string databaseName,
         AzureCosmosDBMongoDBConfig config
     )
     {
-        MongoClientSettings settings = mongoClient.Settings;
         this._config = config;
-        settings.ApplicationName = this._config.ApplicationName;
-        this._mongoClient = new MongoClient(settings);
+        this._mongoClient = mongoClient;
         this._mongoDatabase = this._mongoClient.GetDatabase(databaseName);
     }
 
@@ -318,7 +318,10 @@ protected virtual void Dispose(bool disposing)
     {
         if (disposing)
         {
-            this._mongoClient.Cluster.Dispose();
+            if (this._ownsMongoClient)
+            {
+                this._mongoClient.Cluster.Dispose();
+            }
         }
     }
 

dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTests.cs

@@ -49,7 +49,6 @@ public async Task ItCanBatchUpsertGetRemoveAsync(bool withEmbeddings)
         var memoryStore = this._fixture.MemoryStore;
         var records = DataHelper.CreateBatchRecords(Count);
 
-        await memoryStore.CreateCollectionAsync(collectionName);
         var keys = await memoryStore.UpsertBatchAsync(collectionName, records).ToListAsync();
         var actualRecords = await memoryStore
             .GetBatchAsync(collectionName, keys, withEmbeddings: withEmbeddings)
@@ -86,6 +85,12 @@ public async Task ItCanGetNearestMatchesAsync(int limit, bool withEmbeddings)
         var memoryStore = this._fixture.MemoryStore;
         var searchEmbedding = DataHelper.VectorSearchTestEmbedding;
         var nearestMatchesExpected = DataHelper.VectorSearchExpectedResults;
+        var records = DataHelper.VectorSearchTestRecords;
+
+        var keys = await memoryStore.UpsertBatchAsync(collectionName, records).ToListAsync();
+        var actualRecords = await memoryStore
+            .GetBatchAsync(collectionName, keys, withEmbeddings: withEmbeddings)
+            .ToListAsync();
 
         var nearestMatchesActual = await memoryStore
             .GetNearestMatchesAsync(

dotnet/src/IntegrationTests/Connectors/Memory/AzureCosmosDBMongoDB/AzureCosmosDBMongoDBMemoryStoreTestsFixture.cs

@@ -28,7 +28,6 @@ public AzureCosmosDBMongoDBMemoryStoreTestsFixture()
             )
             .AddEnvironmentVariables()
             .Build();
-
         var connectionString = GetSetting(configuration, "ConnectionString");
         this.DatabaseName = "DotNetSKTestDB";
         this.CollectionName = "DotNetSKTestCollection";
@@ -42,7 +41,6 @@ public AzureCosmosDBMongoDBMemoryStoreTestsFixture()
     public async Task InitializeAsync()
     {
         await this.MemoryStore.CreateCollectionAsync(this.CollectionName);
-
         await this
             .MemoryStore.UpsertBatchAsync(this.CollectionName, DataHelper.VectorSearchTestRecords)
             .ToListAsync();