diff --git a/tests/Integration/Momento.Sdk.Tests/Utils.cs b/tests/Integration/Momento.Sdk.Tests/Utils.cs index 57b8689f..aa29b081 100644 --- a/tests/Integration/Momento.Sdk.Tests/Utils.cs +++ b/tests/Integration/Momento.Sdk.Tests/Utils.cs @@ -1,3 +1,6 @@ +using System; +using Momento.Sdk.Requests.Vector; +using Momento.Sdk.Responses.Vector; namespace Momento.Sdk.Tests.Integration; /// @@ -30,4 +33,49 @@ public static void CreateCacheForTest(ICacheClient cacheClient, string cacheName throw new Exception($"Error when creating cache: {result}"); } } + + public static _WithVectorIndex WithVectorIndex(IPreviewVectorIndexClient vectorIndexClient, IndexInfo indexInfo) + { + return new _WithVectorIndex(vectorIndexClient, indexInfo); + } + + public static _WithVectorIndex WithVectorIndex(IPreviewVectorIndexClient vectorIndexClient, string indexName, int numDimensions, SimilarityMetric similarityMetric = SimilarityMetric.CosineSimilarity) + { + return WithVectorIndex(vectorIndexClient, new IndexInfo(indexName, numDimensions, similarityMetric)); + } + + public class _WithVectorIndex : IDisposable + { + public IPreviewVectorIndexClient VectorIndexClient { get; } + + public IndexInfo IndexInfo { get; } + + public _WithVectorIndex(IPreviewVectorIndexClient vectorIndexClient, IndexInfo indexInfo) + { + VectorIndexClient = vectorIndexClient; + IndexInfo = indexInfo; + + // This usually isn't kosher in a constructor; because we want this class to encapsulate + // index creation and deletion in a using block, the class has to do RAII. + var createResponse = vectorIndexClient.CreateIndexAsync(indexInfo.Name, indexInfo.NumDimensions, indexInfo.SimilarityMetric).Result; + if (createResponse is not (CreateIndexResponse.Success or CreateIndexResponse.AlreadyExists)) + { + throw new Exception($"Error when creating index: {createResponse}"); + } + } + + public _WithVectorIndex(IPreviewVectorIndexClient vectorIndexClient, string indexName, int numDimensions, SimilarityMetric similarityMetric = SimilarityMetric.CosineSimilarity) + : this(vectorIndexClient, new IndexInfo(indexName, numDimensions, similarityMetric)) + { + } + + public void Dispose() + { + var deleteResponse = VectorIndexClient.DeleteIndexAsync(IndexInfo.Name).Result; + if (deleteResponse is not DeleteIndexResponse.Success) + { + throw new Exception($"Error when deleting index: {deleteResponse}"); + } + } + } } diff --git a/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs b/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs index f92266c5..72cb8e1e 100644 --- a/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs +++ b/tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs @@ -32,21 +32,13 @@ public static IEnumerable CreateAndListIndexTestData [MemberData(nameof(CreateAndListIndexTestData))] public async Task CreateListDelete_HappyPath(IndexInfo indexInfo) { - try + using (Utils.WithVectorIndex(vectorIndexClient, indexInfo)) { - var createResponse = await vectorIndexClient.CreateIndexAsync(indexInfo.Name, indexInfo.NumDimensions, indexInfo.SimilarityMetric); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - var listResponse = await vectorIndexClient.ListIndexesAsync(); Assert.True(listResponse is ListIndexesResponse.Success, $"Unexpected response: {listResponse}"); var listOk = (ListIndexesResponse.Success)listResponse; Assert.Contains(indexInfo, listOk.Indexes); } - finally - { - var deleteResponse = await vectorIndexClient.DeleteIndexAsync(indexInfo.Name); - Assert.True(deleteResponse is DeleteIndexResponse.Success, $"Unexpected response: {deleteResponse}"); - } } [Fact] @@ -54,12 +46,11 @@ public async Task CreateIndexAsync_AlreadyExistsError() { var indexName = Utils.TestVectorIndexName(); const int numDimensions = 3; - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, numDimensions); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - var createAgainResponse = await vectorIndexClient.CreateIndexAsync(indexName, numDimensions); - Assert.True(createAgainResponse is CreateIndexResponse.AlreadyExists, $"Unexpected response: {createAgainResponse}"); + using (Utils.WithVectorIndex(vectorIndexClient, indexName, numDimensions)) + { + var createAgainResponse = await vectorIndexClient.CreateIndexAsync(indexName, numDimensions); + Assert.True(createAgainResponse is CreateIndexResponse.AlreadyExists, $"Unexpected response: {createAgainResponse}"); + } } [Fact] diff --git a/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs b/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs index da0e9f9b..a24e4eb4 100644 --- a/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs +++ b/tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs @@ -110,11 +110,7 @@ public async Task UpsertAndSearch_InnerProduct(SearchDelegate searchDelega AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2, SimilarityMetric.InnerProduct); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct)) { var items = new List { @@ -134,10 +130,6 @@ public async Task UpsertAndSearch_InnerProduct(SearchDelegate searchDelega new("test_item", 5.0f) }, items.Select(i => i.Vector).ToList()); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } [Theory] @@ -146,11 +138,7 @@ public async Task UpsertAndSearch_CosineSimilarity(SearchDelegate searchDe AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2)) { var items = new List { @@ -173,10 +161,6 @@ public async Task UpsertAndSearch_CosineSimilarity(SearchDelegate searchDe new("test_item_3", -1.0f) }, items.Select(i => i.Vector).ToList()); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } [Theory] @@ -185,12 +169,7 @@ public async Task UpsertAndSearch_EuclideanSimilarity(SearchDelegate searc AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = - await vectorIndexClient.CreateIndexAsync(indexName, 2, SimilarityMetric.EuclideanSimilarity); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.EuclideanSimilarity)) { var items = new List { @@ -213,10 +192,6 @@ public async Task UpsertAndSearch_EuclideanSimilarity(SearchDelegate searc new("test_item_3", 8.0f) }, items.Select(i => i.Vector).ToList()); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } [Theory] @@ -225,11 +200,7 @@ public async Task UpsertAndSearch_TopKLimit(SearchDelegate searchDelegate, AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2, SimilarityMetric.InnerProduct); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct)) { var items = new List { @@ -255,10 +226,6 @@ public async Task UpsertAndSearch_TopKLimit(SearchDelegate searchDelegate, new() { 3.0f, 4.0f } }); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } [Theory] @@ -267,11 +234,7 @@ public async Task UpsertAndSearch_WithMetadata(SearchDelegate searchDelega AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2, SimilarityMetric.InnerProduct); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct)) { var items = new List { @@ -335,10 +298,6 @@ public async Task UpsertAndSearch_WithMetadata(SearchDelegate searchDelega new Dictionary { { "key1", "value1" } }) }, expectedVectors); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } [Theory] @@ -347,11 +306,7 @@ public async Task UpsertAndSearch_WithDiverseMetadata(SearchDelegate searc AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2, SimilarityMetric.InnerProduct); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, SimilarityMetric.InnerProduct)) { var metadata = new Dictionary { @@ -381,10 +336,6 @@ public async Task UpsertAndSearch_WithDiverseMetadata(SearchDelegate searc new("test_item_1", 5.0f, metadata) }, items.Select(i => i.Vector).ToList()); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } public static IEnumerable SearchThresholdTestCases => @@ -423,11 +374,7 @@ public async Task Search_PruneBasedOnThreshold(SimilarityMetric similarityMet List thresholds, SearchDelegate searchDelegate, AssertOnSearchResponse assertOnSearchResponse) { var indexName = Utils.TestVectorIndexName(); - - var createResponse = await vectorIndexClient.CreateIndexAsync(indexName, 2, similarityMetric); - Assert.True(createResponse is CreateIndexResponse.Success, $"Unexpected response: {createResponse}"); - - try + using (Utils.WithVectorIndex(vectorIndexClient, indexName, 2, similarityMetric)) { var items = new List { @@ -467,9 +414,5 @@ public async Task Search_PruneBasedOnThreshold(SimilarityMetric similarityMet await searchDelegate.Invoke(vectorIndexClient, indexName, queryVector, 3, scoreThreshold: thresholds[2]); assertOnSearchResponse.Invoke(searchResponse, new List(), new List>()); } - finally - { - await vectorIndexClient.DeleteIndexAsync(indexName); - } } }