Skip to content

Commit 31b9e5d

Browse files
jairad26jairad26
committed
[TST] add proptest for config & schema reconciliation
1 parent a021c46 commit 31b9e5d

File tree

5 files changed

+1292
-0
lines changed

5 files changed

+1292
-0
lines changed

‎rust/types/src/collection_configuration.rs‎

Lines changed: 357 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1079,4 +1079,361 @@ mod tests {
10791079
))
10801080
);
10811081
}
1082+
1083+
#[cfg(feature = "testing")]
1084+
mod proptests {
1085+
use super::*;
1086+
use crate::hnsw_configuration::Space;
1087+
use crate::proptest_utils::strategies::*;
1088+
use crate::{HnswConfiguration, MetadataValue, SpannConfiguration};
1089+
use proptest::prelude::*;
1090+
use proptest::test_runner::TestCaseError;
1091+
1092+
fn space_to_metadata_str(space: &Space) -> &'static str {
1093+
match space {
1094+
Space::L2 => "l2",
1095+
Space::Cosine => "cosine",
1096+
Space::Ip => "ip",
1097+
}
1098+
}
1099+
1100+
fn metadata_from_hnsw_config(config: &InternalHnswConfiguration) -> Metadata {
1101+
let mut metadata = Metadata::new();
1102+
metadata.insert(
1103+
"hnsw:space".to_string(),
1104+
MetadataValue::Str(space_to_metadata_str(&config.space).to_string()),
1105+
);
1106+
metadata.insert(
1107+
"hnsw:construction_ef".to_string(),
1108+
MetadataValue::Int(config.ef_construction as i64),
1109+
);
1110+
metadata.insert(
1111+
"hnsw:search_ef".to_string(),
1112+
MetadataValue::Int(config.ef_search as i64),
1113+
);
1114+
metadata.insert(
1115+
"hnsw:M".to_string(),
1116+
MetadataValue::Int(config.max_neighbors as i64),
1117+
);
1118+
metadata.insert(
1119+
"hnsw:num_threads".to_string(),
1120+
MetadataValue::Int(config.num_threads as i64),
1121+
);
1122+
metadata.insert(
1123+
"hnsw:resize_factor".to_string(),
1124+
MetadataValue::Float(config.resize_factor),
1125+
);
1126+
metadata.insert(
1127+
"hnsw:sync_threshold".to_string(),
1128+
MetadataValue::Int(config.sync_threshold as i64),
1129+
);
1130+
metadata.insert(
1131+
"hnsw:batch_size".to_string(),
1132+
MetadataValue::Int(config.batch_size as i64),
1133+
);
1134+
metadata
1135+
}
1136+
1137+
fn metadata_hnsw_strategy() -> impl Strategy<Value = (InternalHnswConfiguration, Metadata)>
1138+
{
1139+
internal_hnsw_configuration_strategy().prop_map(|config| {
1140+
let metadata = metadata_from_hnsw_config(&config);
1141+
(config, metadata)
1142+
})
1143+
}
1144+
1145+
fn assert_spann_is_default_with_space(
1146+
config: &InternalSpannConfiguration,
1147+
expected_space: Space,
1148+
) -> Result<(), TestCaseError> {
1149+
let default_config = InternalSpannConfiguration {
1150+
space: expected_space,
1151+
..InternalSpannConfiguration::default()
1152+
};
1153+
prop_assert_eq!(config, &default_config);
1154+
Ok(())
1155+
}
1156+
1157+
fn assert_hnsw_is_default_with_space(
1158+
config: &InternalHnswConfiguration,
1159+
expected_space: Space,
1160+
) -> Result<(), TestCaseError> {
1161+
let default_config = InternalHnswConfiguration {
1162+
space: expected_space,
1163+
..InternalHnswConfiguration::default()
1164+
};
1165+
prop_assert_eq!(config, &default_config);
1166+
Ok(())
1167+
}
1168+
1169+
proptest! {
1170+
#[test]
1171+
fn try_from_config_roundtrip_internal(
1172+
internal_config in internal_collection_configuration_strategy()
1173+
) {
1174+
let collection_config: CollectionConfiguration = internal_config.clone().into();
1175+
let default_knn = match &internal_config.vector_index {
1176+
VectorIndexConfiguration::Hnsw(_) => KnnIndex::Hnsw,
1177+
VectorIndexConfiguration::Spann(_) => KnnIndex::Spann,
1178+
};
1179+
1180+
let result = InternalCollectionConfiguration::try_from_config(
1181+
collection_config.clone(),
1182+
default_knn,
1183+
None,
1184+
)
1185+
.expect("conversion should succeed");
1186+
1187+
let embedding_function = internal_config.embedding_function.clone();
1188+
let expected_vector_index = match &internal_config.vector_index {
1189+
VectorIndexConfiguration::Hnsw(original) => {
1190+
let external: HnswConfiguration = original.clone().into();
1191+
let expected_internal: InternalHnswConfiguration = external.clone().into();
1192+
match &result.vector_index {
1193+
VectorIndexConfiguration::Hnsw(converted) => {
1194+
prop_assert_eq!(converted, &expected_internal);
1195+
}
1196+
_ => prop_assert!(false, "expected HNSW configuration"),
1197+
}
1198+
VectorIndexConfiguration::Hnsw(expected_internal)
1199+
}
1200+
VectorIndexConfiguration::Spann(original) => {
1201+
let external: SpannConfiguration = original.clone().into();
1202+
let expected_internal: InternalSpannConfiguration = external.clone().into();
1203+
match &result.vector_index {
1204+
VectorIndexConfiguration::Spann(converted) => {
1205+
prop_assert_eq!(converted, &expected_internal);
1206+
}
1207+
_ => prop_assert!(false, "expected SPANN configuration"),
1208+
}
1209+
VectorIndexConfiguration::Spann(expected_internal)
1210+
}
1211+
};
1212+
1213+
prop_assert_eq!(
1214+
result.embedding_function.clone(),
1215+
embedding_function.clone()
1216+
);
1217+
let expected = InternalCollectionConfiguration {
1218+
vector_index: expected_vector_index,
1219+
embedding_function: embedding_function.clone(),
1220+
};
1221+
prop_assert_eq!(result, expected);
1222+
1223+
let opposite_knn = match &internal_config.vector_index {
1224+
VectorIndexConfiguration::Hnsw(_) => KnnIndex::Spann,
1225+
VectorIndexConfiguration::Spann(_) => KnnIndex::Hnsw,
1226+
};
1227+
let opposite_result = InternalCollectionConfiguration::try_from_config(
1228+
collection_config,
1229+
opposite_knn,
1230+
None,
1231+
)
1232+
.expect("conversion for opposite default should succeed");
1233+
1234+
prop_assert_eq!(
1235+
opposite_result.embedding_function.clone(),
1236+
internal_config.embedding_function.clone()
1237+
);
1238+
1239+
match (&internal_config.vector_index, &opposite_result.vector_index) {
1240+
(VectorIndexConfiguration::Hnsw(original), VectorIndexConfiguration::Spann(spann)) => {
1241+
let expected_space = original.space.clone();
1242+
assert_spann_is_default_with_space(spann, expected_space)?;
1243+
}
1244+
(VectorIndexConfiguration::Spann(original), VectorIndexConfiguration::Hnsw(hnsw)) => {
1245+
let expected_space = original.space.clone();
1246+
assert_hnsw_is_default_with_space(hnsw, expected_space)?;
1247+
}
1248+
_ => prop_assert!(false, "unexpected opposite conversion result"),
1249+
}
1250+
}
1251+
}
1252+
1253+
proptest! {
1254+
#[test]
1255+
fn try_from_config_uses_metadata_when_configs_absent(
1256+
(expected_hnsw, metadata) in metadata_hnsw_strategy(),
1257+
embedding in embedding_function_strategy(),
1258+
knn in knn_index_strategy(),
1259+
) {
1260+
let collection_config = CollectionConfiguration {
1261+
hnsw: None,
1262+
spann: None,
1263+
embedding_function: embedding.clone(),
1264+
};
1265+
1266+
let result = InternalCollectionConfiguration::try_from_config(
1267+
collection_config,
1268+
knn,
1269+
Some(metadata.clone()),
1270+
)
1271+
.expect("conversion should succeed");
1272+
1273+
match (knn, &result.vector_index) {
1274+
(KnnIndex::Hnsw, VectorIndexConfiguration::Hnsw(hnsw)) => {
1275+
prop_assert_eq!(hnsw, &expected_hnsw);
1276+
}
1277+
(KnnIndex::Spann, VectorIndexConfiguration::Spann(spann)) => {
1278+
prop_assert_eq!(spann.space.clone(), expected_hnsw.space.clone());
1279+
assert_spann_is_default_with_space(spann, expected_hnsw.space.clone())?;
1280+
}
1281+
_ => prop_assert!(false, "unexpected vector index variant"),
1282+
}
1283+
prop_assert_eq!(result.embedding_function.clone(), embedding);
1284+
}
1285+
}
1286+
1287+
proptest! {
1288+
#[test]
1289+
fn try_from_config_uses_metadata_when_hnsw_config_is_default_values(
1290+
(expected_hnsw, metadata) in metadata_hnsw_strategy(),
1291+
embedding in embedding_function_strategy(),
1292+
) {
1293+
let collection_config = CollectionConfiguration {
1294+
hnsw: Some(HnswConfiguration::default()),
1295+
spann: None,
1296+
embedding_function: embedding.clone(),
1297+
};
1298+
1299+
let result = InternalCollectionConfiguration::try_from_config(
1300+
collection_config,
1301+
KnnIndex::Hnsw,
1302+
Some(metadata.clone()),
1303+
)
1304+
.expect("conversion should succeed");
1305+
1306+
match &result.vector_index {
1307+
VectorIndexConfiguration::Hnsw(hnsw) => {
1308+
prop_assert_eq!(hnsw, &expected_hnsw);
1309+
}
1310+
_ => prop_assert!(false, "expected hnsw configuration"),
1311+
}
1312+
prop_assert_eq!(result.embedding_function.clone(), embedding);
1313+
}
1314+
}
1315+
1316+
proptest! {
1317+
#[test]
1318+
fn try_from_config_prefers_spann_when_default_is_spann(
1319+
hnsw_config in internal_hnsw_configuration_strategy(),
1320+
embedding in embedding_function_strategy(),
1321+
) {
1322+
let embedding_clone = embedding.clone();
1323+
let collection_config = CollectionConfiguration {
1324+
hnsw: Some(hnsw_config.clone().into()),
1325+
spann: None,
1326+
embedding_function: embedding,
1327+
};
1328+
1329+
let result = InternalCollectionConfiguration::try_from_config(
1330+
collection_config,
1331+
KnnIndex::Spann,
1332+
None,
1333+
)
1334+
.expect("conversion should succeed");
1335+
1336+
let expected_space = hnsw_config.space.clone();
1337+
match &result.vector_index {
1338+
VectorIndexConfiguration::Spann(spann) => {
1339+
prop_assert_eq!(spann.space.clone(), expected_space.clone());
1340+
assert_spann_is_default_with_space(spann, expected_space)?;
1341+
}
1342+
_ => prop_assert!(false, "expected spann configuration"),
1343+
}
1344+
prop_assert_eq!(result.embedding_function.clone(), embedding_clone);
1345+
}
1346+
}
1347+
1348+
proptest! {
1349+
#[test]
1350+
fn try_from_config_prefers_hnsw_when_default_is_hnsw(
1351+
spann_config in internal_spann_configuration_strategy(),
1352+
embedding in embedding_function_strategy(),
1353+
) {
1354+
let embedding_clone = embedding.clone();
1355+
let collection_config = CollectionConfiguration {
1356+
hnsw: None,
1357+
spann: Some(spann_config.clone().into()),
1358+
embedding_function: embedding,
1359+
};
1360+
1361+
let result = InternalCollectionConfiguration::try_from_config(
1362+
collection_config,
1363+
KnnIndex::Hnsw,
1364+
None,
1365+
)
1366+
.expect("conversion should succeed");
1367+
1368+
let expected_space = spann_config.space.clone();
1369+
match &result.vector_index {
1370+
VectorIndexConfiguration::Hnsw(hnsw) => {
1371+
prop_assert_eq!(hnsw.space.clone(), expected_space.clone());
1372+
assert_hnsw_is_default_with_space(hnsw, expected_space)?;
1373+
}
1374+
_ => prop_assert!(false, "expected hnsw configuration"),
1375+
}
1376+
prop_assert_eq!(result.embedding_function.clone(), embedding_clone);
1377+
}
1378+
}
1379+
1380+
proptest! {
1381+
#[test]
1382+
fn try_from_config_defaults_when_configs_absent(
1383+
embedding in embedding_function_strategy(),
1384+
knn in knn_index_strategy(),
1385+
) {
1386+
let collection_config = CollectionConfiguration {
1387+
hnsw: None,
1388+
spann: None,
1389+
embedding_function: embedding.clone(),
1390+
};
1391+
1392+
let result = InternalCollectionConfiguration::try_from_config(
1393+
collection_config,
1394+
knn,
1395+
None,
1396+
)
1397+
.expect("conversion should succeed");
1398+
1399+
match (knn, &result.vector_index) {
1400+
(KnnIndex::Hnsw, VectorIndexConfiguration::Hnsw(hnsw)) => {
1401+
prop_assert_eq!(hnsw, &InternalHnswConfiguration::default());
1402+
}
1403+
(KnnIndex::Spann, VectorIndexConfiguration::Spann(spann)) => {
1404+
prop_assert_eq!(spann, &InternalSpannConfiguration::default());
1405+
}
1406+
_ => prop_assert!(false, "unexpected vector index variant"),
1407+
}
1408+
prop_assert_eq!(result.embedding_function.clone(), embedding);
1409+
}
1410+
}
1411+
1412+
proptest! {
1413+
#[test]
1414+
fn try_from_config_errors_on_multiple_configs(
1415+
hnsw_config in internal_hnsw_configuration_strategy(),
1416+
spann_config in internal_spann_configuration_strategy(),
1417+
embedding in embedding_function_strategy(),
1418+
knn in knn_index_strategy(),
1419+
) {
1420+
let collection_config = CollectionConfiguration {
1421+
hnsw: Some(hnsw_config.into()),
1422+
spann: Some(spann_config.into()),
1423+
embedding_function: embedding,
1424+
};
1425+
1426+
let result = InternalCollectionConfiguration::try_from_config(
1427+
collection_config,
1428+
knn,
1429+
None,
1430+
);
1431+
1432+
prop_assert!(matches!(
1433+
result,
1434+
Err(CollectionConfigurationToInternalConfigurationError::MultipleVectorIndexConfigurations)
1435+
));
1436+
}
1437+
}
1438+
}
10821439
}

0 commit comments

Comments
 (0)