// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using Microsoft.Extensions.VectorData;
using VectorData.ConformanceTests.Support;
using Xunit;
namespace VectorData.ConformanceTests.ModelTests;
#pragma warning disable S3218 // Inner class members should not shadow outer class "static" or type members
public abstract class DynamicModelTests<TKey>(DynamicModelTests<TKey>.Fixture fixture) : IAsyncLifetime
where TKey : notnull
{
#region Get
[Theory]
[MemberData(nameof(IncludeVectorsData))]
public virtual async Task GetAsync_single_record(bool includeVectors)
{
var expectedRecord = fixture.TestData[0];
var received = await fixture.Collection.GetAsync(
(TKey)expectedRecord[KeyPropertyName]!,
new() { IncludeVectors = includeVectors });
AssertEquivalent(expectedRecord, received, includeVectors, fixture.TestStore.VectorsComparable);
}
[Theory]
[MemberData(nameof(IncludeVectorsData))]
public virtual async Task GetAsync_multiple_records(bool includeVectors)
{
var expectedRecords = fixture.TestData.Take(2);
var ids = expectedRecords.Select(record => record[KeyPropertyName]!);
var received = await fixture.Collection.GetAsync(ids, new() { IncludeVectors = includeVectors }).ToArrayAsync();
foreach (var record in expectedRecords)
{
AssertEquivalent(
record,
received.Single(r => r[KeyPropertyName]!.Equals(record[KeyPropertyName])),
includeVectors,
fixture.TestStore.VectorsComparable);
}
}
[Fact]
public virtual async Task GetAsync_throws_for_null_key()
{
// Skip this test for value type keys
if (default(TKey) is not null)
{
return;
}
ArgumentNullException ex = await Assert.ThrowsAsync<ArgumentNullException>(() => fixture.Collection.GetAsync((TKey)default!));
Assert.Equal("key", ex.ParamName);
}
[Fact]
public virtual async Task GetAsync_throws_for_null_keys()
{
ArgumentNullException ex = await Assert.ThrowsAsync<ArgumentNullException>(() => fixture.Collection.GetAsync(keys: null!).ToArrayAsync().AsTask());
Assert.Equal("keys", ex.ParamName);
}
[Fact]
public virtual async Task GetAsync_returns_null_for_missing_key()
{
TKey key = fixture.GenerateNextKey<TKey>();
Assert.Null(await fixture.Collection.GetAsync(key));
}
[Fact]
public virtual async Task GetAsync_returns_empty_for_empty_keys()
{
Assert.Empty(await fixture.Collection.GetAsync([]).ToArrayAsync());
}
[Theory]
[MemberData(nameof(IncludeVectorsData))]
public virtual async Task GetAsync_with_filter(bool includeVectors)
{
var expectedRecord = fixture.TestData[0];
var results = await fixture.Collection.GetAsync(
r => (int)r[IntegerPropertyName]! == 1,
top: 2,
new() { IncludeVectors = includeVectors })
.ToListAsync();
var receivedRecord = Assert.Single(results);
AssertEquivalent(expectedRecord, receivedRecord, includeVectors, fixture.TestStore.VectorsComparable);
}
[Fact]
public virtual async Task GetAsync_with_filter_by_true()
{
var count = await fixture.Collection.GetAsync(r => true, top: 100).CountAsync();
Assert.Equal(fixture.TestData.Count, count);
Assert.True(count < 100);
}
[Fact]
public virtual async Task GetAsync_with_filter_and_OrderBy()
{
var ascendingNumbers = fixture.TestData
.Where(r => (int)r[IntegerPropertyName]! > 1)
.OrderBy(r => r[IntegerPropertyName])
.Take(2)
.Select(r => (int)r[IntegerPropertyName]!)
.ToList();
var descendingNumbers = fixture.TestData
.Where(r => (int)r[IntegerPropertyName]! > 1)
.OrderByDescending(r => r[IntegerPropertyName])
.Take(2)
.Select(r => (int)r[IntegerPropertyName]!)
.ToList();
// Make sure the actual results are different for ascending/descending, otherwise the test is meaningless
Assert.NotEqual(ascendingNumbers, descendingNumbers);
// Finally, query once with ascending and once with descending, comparing against the expected results above.
var results = await fixture.Collection.GetAsync(
r => (int)r[IntegerPropertyName]! > 1,
top: 2,
new() { OrderBy = o => o.Ascending(r => r[IntegerPropertyName]) })
.Select(r => (int)r[IntegerPropertyName]!)
.ToListAsync();
Assert.Equal(ascendingNumbers, results);
results = await fixture.Collection.GetAsync(
r => (int)r[IntegerPropertyName]! > 1,
top: 2,
new() { OrderBy = o => o.Descending(r => r[IntegerPropertyName]) })
.Select(r => (int)r[IntegerPropertyName]!)
.ToListAsync();
Assert.Equal(descendingNumbers, results);
}
[Fact]
public virtual async Task GetAsync_with_filter_and_multiple_OrderBys()
{
var ascendingNumbers = fixture.TestData
.OrderByDescending(r => r[StringPropertyName])
.ThenBy(r => r[IntegerPropertyName])
.Take(2)
.Select(r => (int)r[IntegerPropertyName]!)
.ToList();
var descendingNumbers = fixture.TestData
.OrderByDescending(r => r[StringPropertyName])
.ThenByDescending(r => r[IntegerPropertyName])
.Take(2)
.Select(r => (int)r[IntegerPropertyName]!)
.ToList();
// Make sure the actual results are different for ascending/descending, otherwise the test is meaningless
Assert.NotEqual(ascendingNumbers, descendingNumbers);
var results = await fixture.Collection.GetAsync(
r => true,
top: 2,
new() { OrderBy = o => o.Descending(r => r[StringPropertyName]).Ascending(r => r[IntegerPropertyName]) })
.Select(r => (int)r[IntegerPropertyName]!)
.ToListAsync();
Assert.Equal(ascendingNumbers, results);
results = await fixture.Collection.GetAsync(
r => true,
top: 2,
new() { OrderBy = o => o.Descending(r => r[StringPropertyName]).Descending(r => r[IntegerPropertyName]) })
.Select(r => (int)r[IntegerPropertyName]!)
.ToListAsync();
Assert.Equal(descendingNumbers, results);
}
[Fact]
public virtual async Task GetAsync_with_filter_and_OrderBy_and_Skip()
{
var results = await fixture.Collection.GetAsync(
r => (int)r[IntegerPropertyName]! > 1,
top: 2,
new() { OrderBy = o => o.Ascending(r => r[IntegerPropertyName]), Skip = 1 })
.Select(r => (int)r[IntegerPropertyName]!)
.ToListAsync();
Assert.Equal(
fixture.TestData
.Where(r => (int)r[IntegerPropertyName]! > 1)
.OrderBy(r => r[IntegerPropertyName])
.Skip(1)
.Take(2)
.Select(r => (int)r[IntegerPropertyName]!),
results);
}
#endregion Get
#region Upsert
[Fact]
public virtual async Task Insert_single_record()
{
TKey expectedKey = fixture.GenerateNextKey<TKey>();
var inserted = new Dictionary<string, object?>
{
[KeyPropertyName] = expectedKey,
[StringPropertyName] = "some",
[IntegerPropertyName] = 123,
[VectorPropertyName] = new ReadOnlyMemory<float>([10, 0, 0])
};
Assert.Null(await Collection.GetAsync(expectedKey));
await Collection.UpsertAsync(inserted);
var received = await Collection.GetAsync(expectedKey, new() { IncludeVectors = true });
AssertEquivalent(inserted, received, includeVectors: true, fixture.TestStore.VectorsComparable);
}
[Fact]
public virtual async Task Update_single_record()
{
var existingRecord = fixture.TestData[1];
var updated = new Dictionary<string, object?>
{
[KeyPropertyName] = existingRecord[KeyPropertyName],
[StringPropertyName] = "different",
[IntegerPropertyName] = 456,
[VectorPropertyName] = new ReadOnlyMemory<float>(Enumerable.Repeat(0.7f, 3).ToArray())
};
Assert.NotNull(await Collection.GetAsync((TKey)existingRecord[KeyPropertyName]!));
await Collection.UpsertAsync(updated);
var received = await Collection.GetAsync((TKey)existingRecord[KeyPropertyName]!, new() { IncludeVectors = true });
AssertEquivalent(updated, received, includeVectors: true, fixture.TestStore.VectorsComparable);
}
[Fact]
public virtual async Task Insert_multiple_records()
{
Dictionary<string, object?>[] newRecords =
[
new()
{
[KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
[IntegerPropertyName] = 100,
[StringPropertyName] = "New record 1",
[VectorPropertyName] = new ReadOnlyMemory<float>([10, 0, 1])
},
new()
{
[KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
[IntegerPropertyName] = 101,
[StringPropertyName] = "New record 2",
[VectorPropertyName] = new ReadOnlyMemory<float>([10, 0, 2])
},
];
var keys = newRecords.Select(record => record[KeyPropertyName]!).ToArray();
Assert.Empty(await Collection.GetAsync(keys).ToArrayAsync());
await Collection.UpsertAsync(newRecords);
var received = await Collection.GetAsync(keys, new() { IncludeVectors = true }).ToArrayAsync();
Assert.Collection(
received.OrderBy(r => r[IntegerPropertyName]),
r => AssertEquivalent(newRecords[0], r, includeVectors: true, fixture.TestStore.VectorsComparable),
r => AssertEquivalent(newRecords[1], r, includeVectors: true, fixture.TestStore.VectorsComparable));
}
#endregion Upsert
#region Delete
[Fact]
public virtual async Task Delete_single_record()
{
var recordToRemove = fixture.TestData[2];
Assert.NotNull(await fixture.Collection.GetAsync((TKey)recordToRemove[KeyPropertyName]!));
await fixture.Collection.DeleteAsync((TKey)recordToRemove[KeyPropertyName]!);
Assert.Null(await fixture.Collection.GetAsync((TKey)recordToRemove[KeyPropertyName]!));
}
[Fact]
public virtual async Task DeleteAsync_does_nothing_for_non_existing_key()
{
TKey key = fixture.GenerateNextKey<TKey>();
await fixture.Collection.DeleteAsync(key);
Assert.Null(await fixture.Collection.GetAsync(key));
}
#endregion Delete
#region Search
[Theory]
[MemberData(nameof(IncludeVectorsData))]
public virtual async Task SearchAsync(bool includeVectors)
{
var expectedRecord = fixture.TestData[0];
var result = await Collection
.SearchAsync(
expectedRecord[VectorPropertyName]!,
top: 1,
new() { IncludeVectors = includeVectors })
.SingleAsync();
AssertEquivalent(expectedRecord, result.Record, includeVectors, fixture.TestStore.VectorsComparable);
}
[Fact]
public virtual async Task SearchAsync_with_Skip()
{
var result = await Collection
.SearchAsync(
fixture.TestData[0][VectorPropertyName]!,
top: 1,
new() { Skip = 1 })
.SingleAsync();
AssertEquivalent(fixture.TestData[1], result.Record, includeVectors: false, fixture.TestStore.VectorsComparable);
}
[Fact]
public virtual async Task SearchAsync_with_Filter()
{
var result = await Collection
.SearchAsync(
fixture.TestData[0][VectorPropertyName]!,
top: 1,
new() { Filter = r => (int)r[IntegerPropertyName]! == 2 })
.SingleAsync();
AssertEquivalent(fixture.TestData[1], result.Record, includeVectors: false, fixture.TestStore.VectorsComparable);
}
#endregion Search
protected static void AssertEquivalent(Dictionary<string, object?> expected, Dictionary<string, object?>? actual, bool includeVectors, bool compareVectors)
{
Assert.NotNull(actual);
Assert.Equal(expected[KeyPropertyName], actual[KeyPropertyName]);
Assert.Equal(expected[StringPropertyName], actual[StringPropertyName]);
Assert.Equal(expected[IntegerPropertyName], actual[IntegerPropertyName]);
if (includeVectors)
{
Assert.Equal(
((ReadOnlyMemory<float>)expected[VectorPropertyName]!).Length,
((ReadOnlyMemory<float>)actual[VectorPropertyName]!).Length);
if (compareVectors)
{
Assert.Equal(
((ReadOnlyMemory<float>)expected[VectorPropertyName]!).ToArray(),
((ReadOnlyMemory<float>)actual[VectorPropertyName]!).ToArray());
}
}
else
{
Assert.False(actual.ContainsKey(VectorPropertyName));
}
}
public const string KeyPropertyName = "key";
public const string StringPropertyName = "text";
public const string IntegerPropertyName = "integer";
public const string VectorPropertyName = "vector";
protected VectorStoreCollection<object, Dictionary<string, object?>> Collection => fixture.Collection;
public abstract class Fixture : DynamicVectorStoreCollectionFixture<TKey>
{
protected override string CollectionNameBase => nameof(DynamicModelTests<int>);
protected override string KeyPropertyName => DynamicModelTests<TKey>.KeyPropertyName;
protected override VectorStoreCollection<object, Dictionary<string, object?>> GetCollection()
=> TestStore.CreateDynamicCollection(CollectionName, CreateRecordDefinition());
public override VectorStoreCollectionDefinition CreateRecordDefinition()
=> new()
{
Properties =
[
new VectorStoreKeyProperty(KeyPropertyName, typeof(TKey)),
new VectorStoreDataProperty(StringPropertyName, typeof(string)) { IsIndexed = true},
new VectorStoreDataProperty(IntegerPropertyName, typeof(int)) { IsIndexed = true },
new VectorStoreVectorProperty(VectorPropertyName, typeof(ReadOnlyMemory<float>), dimensions: 3)
{
DistanceFunction = DistanceFunction,
IndexKind = IndexKind
}
]
};
protected override List<Dictionary<string, object?>> BuildTestData() =>
[
new()
{
[KeyPropertyName] = GenerateNextKey<TKey>(),
[StringPropertyName] = "foo",
[IntegerPropertyName] = 1,
[VectorPropertyName] = new ReadOnlyMemory<float>([1, 2, 3])
},
new()
{
[KeyPropertyName] = GenerateNextKey<TKey>(),
[StringPropertyName] = "bar",
[IntegerPropertyName] = 2,
[VectorPropertyName] = new ReadOnlyMemory<float>([1, 2, 4])
},
new()
{
[KeyPropertyName] = GenerateNextKey<TKey>(),
[StringPropertyName] = "foo", // identical text as above
[IntegerPropertyName] = 3,
[VectorPropertyName] = new ReadOnlyMemory<float>([1, 2, 5])
}
];
}
public ValueTask InitializeAsync()
=> new(fixture.ReseedAsync());
public ValueTask DisposeAsync()
{
GC.SuppressFinalize(this);
return default;
}
public static readonly TheoryData<bool> IncludeVectorsData = [false, true];
}