use super::rerank_test_support as support;
use std::num::NonZeroU64;
use serde_json::{Value, json};
use support::{Fixture, TestResult, patients, request, values};
const FOUR: NonZeroU64 = NonZeroU64::MIN.saturating_add(3);
const SEARCH: &str = "similar(.v; record(\"p0\"); dot)";
fn fixture() -> Result<Fixture, Box<dyn std::error::Error>> {
let mut fixture = Fixture::new(None, &patients())?;
fixture.vectors("v", &[(0, [1.0, 0.0]), (1, [9.0, 0.0]), (2, [3.0, 0.0])])?;
Ok(fixture)
}
fn ids(answer: &Value) -> Result<Vec<&str>, Box<dyn std::error::Error>> {
answer["rows"]
.as_array()
.ok_or("missing rows")?
.iter()
.map(|row| {
row["origin"]["id"]
.as_str()
.ok_or_else(|| "missing id".into())
})
.collect()
}
#[test]
fn similar_after_rg_scores_only_matching_records_and_binds_score() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | rg(\"4B\") | {SEARCH} | limit(2) | '$score'");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p2", "p0"]);
assert_eq!(values(&answer)?, [json!(3.0), json!(1.0)]);
assert_eq!(answer["rows"][0]["origin"]["score"], 3.0);
assert_eq!(answer["rows"][1]["origin"]["score"], 1.0);
assert_eq!(answer["records_read"], 6);
assert_eq!(answer["unscored"], 1);
assert!(answer["next"].is_null());
Ok(())
}
#[test]
fn similar_after_quoted_select_scores_only_its_whole_records() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | 'select(.ward == \"4B\")' | {SEARCH} | limit(2)");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p2", "p0"]);
assert_eq!(
values(&answer)?,
[patients()[2].clone(), patients()[0].clone()]
);
assert_eq!(answer["rows"][0]["origin"]["score"], 3.0);
assert_eq!(answer["rows"][1]["origin"]["score"], 1.0);
assert_eq!(answer["records_read"], 6);
assert!(answer["next"].is_null());
Ok(())
}
fn refused_after(
stage: &str,
name: &str,
) -> Result<crate::serve::storage::Failure, Box<dyn std::error::Error>> {
let fixture = fixture()?;
let query = format!("patients | {stage} | {SEARCH}");
let error = fixture
.run(&request(&query, FOUR), &[])
.err()
.ok_or("incomplete records searched")?;
assert_eq!(error.code, "query_stage_not_available");
assert!(
error.message.contains(&format!(
"similar needs whole records; {name} before it does not keep them"
)),
"{}",
error.message
);
Ok(error)
}
#[test]
fn similar_after_limit_is_refused_by_name() -> TestResult {
let error = refused_after("limit(2)", "limit")?;
assert_eq!(error.code, "query_stage_not_available");
Ok(())
}
#[test]
fn similar_after_count_is_refused_by_name() -> TestResult {
let error = refused_after("count", "count")?;
assert_eq!(error.code, "query_stage_not_available");
Ok(())
}
#[test]
fn similar_after_group_by_is_refused_by_name() -> TestResult {
let error = refused_after("group_by(.ward)", "group_by")?;
assert_eq!(error.code, "query_stage_not_available");
Ok(())
}
#[test]
fn similar_after_percentile_is_refused_by_name() -> TestResult {
let error = refused_after("percentile(.admitted; 95)", "percentile")?;
assert_eq!(error.code, "query_stage_not_available");
Ok(())
}
#[test]
fn similar_after_rerank_is_refused_by_name() -> TestResult {
let error = refused_after("rerank(.v; record(\"p0\"); dot)", "rerank")?;
assert_eq!(error.code, "query_stage_not_available");
Ok(())
}
#[test]
fn similar_keeps_and_counts_records_without_vectors_last() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | rg(\"4B\") | {SEARCH}");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p2", "p0", "p3"]);
assert_eq!(answer["unscored"], 1);
assert!(answer["rows"][2]["origin"]["score"].is_null());
assert_eq!(values(&answer)?[2], patients()[3]);
Ok(())
}
#[test]
fn similar_whole_drain_refuses_the_records_read_limit() -> TestResult {
let mut fixture = fixture()?;
fixture.scan(NonZeroU64::MIN.saturating_add(2));
let query = format!("patients | rg(\"4B\") | {SEARCH} | limit(1)");
let error = fixture
.run(&request(&query, FOUR), &[])
.err()
.ok_or("drain was silently truncated")?;
assert_eq!(error.code, "query_scan_limit");
assert!(error.message.contains('3'), "{}", error.message);
Ok(())
}
#[test]
fn similar_vector_reads_share_the_records_read_limit() -> TestResult {
let mut fixture = fixture()?;
fixture.scan(NonZeroU64::MIN.saturating_add(4));
let query = format!("patients | rg(\"4B\") | {SEARCH} | limit(1)");
let error = fixture
.run(&request(&query, FOUR), &[])
.err()
.ok_or("vector reads exceeded allowance")?;
assert_eq!(error.code, "query_scan_limit");
assert!(error.message.contains('5'), "{}", error.message);
Ok(())
}
#[test]
fn similar_after_scalar_jq_uses_the_existing_missing_record_refusal() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | '.name' | {SEARCH}");
let error = fixture
.run(&request(&query, FOUR), &[])
.err()
.ok_or("scalar jq was searched by origin")?;
assert_eq!(error.code, "query_search_record_missing");
assert_eq!(
error.message,
"the record named has no vector in .v that this request may read"
);
Ok(())
}
#[test]
fn similar_equal_scores_are_ordered_by_record_id() -> TestResult {
let mut fixture = Fixture::new(None, &patients())?;
fixture.vectors("v", &[(0, [1.0, 0.0]), (1, [9.0, 0.0]), (2, [1.0, 0.0])])?;
let query = format!("patients | sort(.admitted; desc) | {SEARCH} | limit(3)");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p1", "p0", "p2"]);
assert_eq!(answer["rows"][1]["origin"]["score"], 1.0);
assert_eq!(answer["rows"][2]["origin"]["score"], 1.0);
let first_query = format!("patients | {SEARCH} | limit(3)");
let first = fixture.run(&request(&first_query, FOUR), &[])?;
assert_eq!(ids(&first)?, ids(&answer)?);
Ok(())
}
#[test]
fn similar_whole_drain_refuses_its_memory_limit() -> TestResult {
let mut fixture = fixture()?;
fixture.memory(NonZeroU64::MIN);
let query = format!("patients | 'select(.ward == \"4B\")' | {SEARCH} | limit(1)");
let error = fixture
.run(&request(&query, FOUR), &[])
.err()
.ok_or("drain exceeded query memory")?;
assert_eq!(error.code, "query_memory_limit");
assert!(error.message.contains("similar"), "{}", error.message);
Ok(())
}
#[test]
fn similar_after_grep_sort_and_shape_keeps_filtered_origins() -> TestResult {
let fixture = fixture()?;
let query = format!(
"patients | grep(\"4B\") | sort(.admitted) | {{ name ward }} | {SEARCH} | limit(2)"
);
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p2", "p0"]);
assert_eq!(
values(&answer)?,
[
json!({"name":"Cy", "ward":"4B"}),
json!({"name":"Ada", "ward":"4B"})
]
);
assert_eq!(answer["rows"][0]["origin"]["score"], 3.0);
Ok(())
}
#[test]
fn similar_after_a_single_collection_root_shape_keeps_origins() -> TestResult {
let fixture = fixture()?;
let query = format!("{{ patients {{ name ward }} }} | {SEARCH} | limit(2)");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p1", "p2"]);
assert_eq!(answer["rows"][0]["origin"]["score"], 9.0);
Ok(())
}
#[test]
fn similar_refuses_an_output_page_cut_instead_of_returning_next() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | rg(\"4B\") | {SEARCH}");
let error = fixture
.run(&request(&query, NonZeroU64::MIN), &[])
.err()
.ok_or("partial similarity escaped")?;
assert_eq!(error.code, "query_output_limit");
Ok(())
}
#[test]
fn similar_after_filter_cannot_read_a_denied_question() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | rg(\"4B\") | {SEARCH}");
let error = fixture
.run(&request(&query, FOUR), &[("patients", "p0")])
.err()
.ok_or("denied question was read")?;
assert_eq!(error.code, "query_search_record_missing");
Ok(())
}
#[test]
fn similar_after_labelled_match_keeps_only_authorized_candidates() -> TestResult {
let fixture = fixture()?;
let query = format!("match (p:patients) | {SEARCH} | limit(2)");
let answer = fixture.run(&request(&query, FOUR), &[("patients", "p1")])?;
assert_eq!(ids(&answer)?, ["p2", "p0"]);
assert_eq!(answer["rows"][0]["origin"]["score"], 3.0);
assert_eq!(values(&answer)?[0], json!({"p": patients()[2]}));
Ok(())
}
#[test]
fn similar_after_window_keeps_the_records_in_that_window() -> TestResult {
let fixture = fixture()?;
let query = format!("patients | window(0; $now) | {SEARCH} | limit(2)");
let answer = fixture.run(&request(&query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p1", "p2"]);
assert_eq!(answer["rows"][0]["origin"]["score"], 9.0);
Ok(())
}
#[test]
fn similar_l2_after_filter_orders_the_nearest_vector_first() -> TestResult {
let fixture = fixture()?;
let answer = fixture.run(
&request(
"patients | rg(\"4B\") | similar(.v; record(\"p0\"); l2)",
FOUR,
),
&[],
)?;
assert_eq!(ids(&answer)?, ["p0", "p2", "p3"]);
assert_eq!(answer["rows"][0]["origin"]["score"], 0.0);
assert_eq!(answer["unscored"], 1);
Ok(())
}