Skip to content

Commit

Permalink
fix: clean up vector indexes in integration tests (#521)
Browse files Browse the repository at this point in the history 
Refactor integration tests to use Disposable pattern throughout. A helper class creates an index on initialization and cleans up on disposal. Integration test authors then wrap code that needs an index in a using block. Upon exit of the using block, the index will be deleted.
  • Loading branch information
@malandis
malandis authored Nov 16, 2023
1 parent 041d537 commit 3f17734
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 79 deletions.
48 changes: 48 additions & 0 deletions tests/Integration/Momento.Sdk.Tests/Utils.cs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
using System;
using Momento.Sdk.Requests.Vector;
using Momento.Sdk.Responses.Vector;
namespace Momento.Sdk.Tests.Integration;

/// <summary>
Expand Down Expand Up @@ -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}");
}
}
}
}
21 changes: 6 additions & 15 deletions tests/Integration/Momento.Sdk.Tests/VectorIndexControlTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,34 +32,25 @@ public static IEnumerable<object[]> 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]
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]
Expand Down
71 changes: 7 additions & 64 deletions tests/Integration/Momento.Sdk.Tests/VectorIndexDataTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,7 @@ public async Task UpsertAndSearch_InnerProduct<T>(SearchDelegate<T> searchDelega
AssertOnSearchResponse<T> 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<Item>
{
Expand All @@ -134,10 +130,6 @@ public async Task UpsertAndSearch_InnerProduct<T>(SearchDelegate<T> searchDelega
new("test_item", 5.0f)
}, items.Select(i => i.Vector).ToList());
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

[Theory]
Expand All @@ -146,11 +138,7 @@ public async Task UpsertAndSearch_CosineSimilarity<T>(SearchDelegate<T> searchDe
AssertOnSearchResponse<T> 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<Item>
{
Expand All @@ -173,10 +161,6 @@ public async Task UpsertAndSearch_CosineSimilarity<T>(SearchDelegate<T> searchDe
new("test_item_3", -1.0f)
}, items.Select(i => i.Vector).ToList());
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

[Theory]
Expand All @@ -185,12 +169,7 @@ public async Task UpsertAndSearch_EuclideanSimilarity<T>(SearchDelegate<T> searc
AssertOnSearchResponse<T> 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<Item>
{
Expand All @@ -213,10 +192,6 @@ public async Task UpsertAndSearch_EuclideanSimilarity<T>(SearchDelegate<T> searc
new("test_item_3", 8.0f)
}, items.Select(i => i.Vector).ToList());
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

[Theory]
Expand All @@ -225,11 +200,7 @@ public async Task UpsertAndSearch_TopKLimit<T>(SearchDelegate<T> searchDelegate,
AssertOnSearchResponse<T> 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<Item>
{
Expand All @@ -255,10 +226,6 @@ public async Task UpsertAndSearch_TopKLimit<T>(SearchDelegate<T> searchDelegate,
new() { 3.0f, 4.0f }
});
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

[Theory]
Expand All @@ -267,11 +234,7 @@ public async Task UpsertAndSearch_WithMetadata<T>(SearchDelegate<T> searchDelega
AssertOnSearchResponse<T> 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<Item>
{
Expand Down Expand Up @@ -335,10 +298,6 @@ public async Task UpsertAndSearch_WithMetadata<T>(SearchDelegate<T> searchDelega
new Dictionary<string, MetadataValue> { { "key1", "value1" } })
}, expectedVectors);
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

[Theory]
Expand All @@ -347,11 +306,7 @@ public async Task UpsertAndSearch_WithDiverseMetadata<T>(SearchDelegate<T> searc
AssertOnSearchResponse<T> 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<string, MetadataValue>
{
Expand Down Expand Up @@ -381,10 +336,6 @@ public async Task UpsertAndSearch_WithDiverseMetadata<T>(SearchDelegate<T> searc
new("test_item_1", 5.0f, metadata)
}, items.Select(i => i.Vector).ToList());
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}

public static IEnumerable<object[]> SearchThresholdTestCases =>
Expand Down Expand Up @@ -423,11 +374,7 @@ public async Task Search_PruneBasedOnThreshold<T>(SimilarityMetric similarityMet
List<float> thresholds, SearchDelegate<T> searchDelegate, AssertOnSearchResponse<T> 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<Item>
{
Expand Down Expand Up @@ -467,9 +414,5 @@ public async Task Search_PruneBasedOnThreshold<T>(SimilarityMetric similarityMet
await searchDelegate.Invoke(vectorIndexClient, indexName, queryVector, 3, scoreThreshold: thresholds[2]);
assertOnSearchResponse.Invoke(searchResponse, new List<SearchHit>(), new List<List<float>>());
}
finally
{
await vectorIndexClient.DeleteIndexAsync(indexName);
}
}
}

0 comments on commit 3f17734

Please sign in to comment.