use std::num::NonZeroU32;
use super::{Argument, Plan, Question, Refused, Source, Span, Stage, pipeline};
use crate::serve::protocol::Metric;
fn read(text: &str) -> Result<Plan, String> {
pipeline(text).map_err(|error| error.to_string())
}
fn collection_jq(text: &str, collection: &str, jq: &str) -> Result<(), String> {
let plan = read(text)?;
assert_eq!(
plan.source,
Source::Collection {
name: collection.to_owned(),
arguments: Vec::new(),
}
);
let values: Vec<_> = plan.stages.iter().map(|stage| &stage.value).collect();
if text.contains('|') {
assert_eq!(values, vec![&Stage::Jq(jq.to_owned())]);
} else {
assert!(values.is_empty());
}
Ok(())
}
#[test]
fn a_pipeline_gives_its_collection_and_its_exact_jq() -> Result<(), String> {
collection_jq("patients", "patients", ".")?;
collection_jq("patients | '.name'", "patients", ".name")?;
collection_jq(
r#"collection("studio.pages") | 'select(.a == "O''Brien")'"#,
"studio.pages",
r#"select(.a == "O'Brien")"#,
)?;
collection_jq("studio.pages", "studio.pages", ".")?;
let alone = read(r#"'select(.status == "open")'"#)?;
assert_eq!(alone.source, Source::All);
assert_eq!(alone.stages.len(), 1);
assert_eq!(
alone.stages[0].value,
Stage::Jq(r#"select(.status == "open")"#.to_owned())
);
Ok(())
}
#[test]
fn later_stages_are_not_yet_and_an_unclosed_quote_is_a_parse_error() -> Result<(), String> {
let not_yet = |text| matches!(pipeline(text), Err(Refused::NotYet(..)));
assert!(not_yet(r#"patients | match "x""#));
assert!(not_yet("patients | '.a' | '.b'"));
let arguments = read(r#"patients(ward: "4B")"#)?;
assert_eq!(
arguments.source,
Source::Collection {
name: "patients".to_owned(),
arguments: vec![Argument {
field: "ward".to_owned(),
value: serde_json::json!("4B"),
}],
}
);
assert!(arguments.stages.is_empty());
let shaped = read("patients | { name }")?;
assert_eq!(shaped.stages.len(), 1);
assert!(matches!(shaped.stages[0].value, Stage::Shape(_)));
assert!(matches!(
pipeline("patients | '.a"),
Err(Refused::Parse(..))
));
Ok(())
}
#[test]
fn plain_jq_is_no_pipeline() {
for text in [".", ".a | .b", "patients | .name"] {
assert!(matches!(pipeline(text), Err(Refused::Parse(..))), "{text}");
}
}
#[test]
fn a_search_by_a_stored_record_is_read_into_the_plan() -> Result<(), String> {
let cases = [
(
r#"notes | similar(.body_vec; record("n-102"); cosine) | limit(10)"#,
Metric::Cosine,
10,
".",
),
(
concat!(
r#"notes | similar(.body_vec; record("n-102"); cosine) | limit(10)"#,
r#" | '{id: $id, score: $score}'"#
),
Metric::Cosine,
10,
"{id: $id, score: $score}",
),
(
r#"notes|similar( .body_vec ; record( "n-102" ) ; l2 )|limit( 3 )"#,
Metric::L2,
3,
".",
),
(
r#"notes | similar(.body_vec; record("n-102"); dot) | limit(1)"#,
Metric::Dot,
1,
".",
),
];
for (text, measure, limit, jq) in cases {
let plan = read(text)?;
let limit = NonZeroU32::new(limit).ok_or("zero limit")?;
assert_eq!(
plan.source,
Source::Collection {
name: "notes".to_owned(),
arguments: Vec::new()
}
);
let mut expected = vec![
Stage::Similar {
field: "body_vec".to_owned(),
question: Question::Record("n-102".to_owned()),
measure,
},
Stage::Limit(limit),
];
if jq != "." {
expected.push(Stage::Jq(jq.to_owned()));
}
assert_eq!(
plan.stages
.iter()
.map(|stage| &stage.value)
.collect::<Vec<_>>(),
expected.iter().collect::<Vec<_>>()
);
}
Ok(())
}
fn refusal(text: &str) -> (&'static str, String) {
match pipeline(text) {
Ok(plan) => ("accepted", format!("{plan:?}")),
Err(refused) => {
let words = refused.to_string();
let message = words
.split_once(" at bytes ")
.map_or(words.as_str(), |(message, _)| message);
(refused.code(), message.to_owned())
}
}
}
#[test]
fn search_forms_not_built_yet_are_refused_by_name() -> Result<(), String> {
let cases = [
(
r#"notes | similar(.body_vec; record("n-102")) | limit(5)"#,
"query_stage_not_available",
"similar names its measure until a field declares one: cosine, l2 or dot",
Span { start: 42, end: 43 },
),
(
r#"similar(.body_vec; record("n-102"); cosine) | limit(5)"#,
"query_parse",
"`similar` is a stage and must follow a collection",
Span { start: 0, end: 7 },
),
(
r#"notes | similar(.a, .b; record("n-102"); cosine) | limit(5)"#,
"query_stage_not_available",
"a search of two fields is not available yet",
Span { start: 18, end: 19 },
),
(
concat!(
r#"notes | similar(.a; record("x"); cosine) | limit(5)"#,
r#" | similar(.b; record("x"); cosine)"#
),
"query_stage_not_available",
"similar needs whole records; limit before it does not keep them",
Span { start: 54, end: 86 },
),
(
"notes | rerank",
"query_parse",
"`rerank` takes its arguments in brackets",
Span { start: 8, end: 14 },
),
(
"notes | sort_by(.a)",
"query_stage_not_available",
"a `sort_by` stage is not available yet",
Span { start: 8, end: 15 },
),
];
for (text, code, words, span) in cases {
assert_eq!(refusal(text), (code, words.to_owned()), "{text}");
let error = pipeline(text).err().ok_or("unserved search was accepted")?;
assert_eq!(error.span(), span, "{text}");
assert_eq!(
error.to_string(),
format!("{words} at bytes {}..{}", span.start, span.end),
"{text}"
);
}
let after_jq =
read(r#"notes | '.a' | similar(.body_vec; record("n-102"); cosine) | limit(5)"#)?;
assert_eq!(after_jq.stages.len(), 3);
assert_eq!(after_jq.stages[0].value, Stage::Jq(".a".to_owned()));
assert_eq!(
after_jq.stages[1].value,
Stage::Similar {
field: "body_vec".to_owned(),
question: Question::Record("n-102".to_owned()),
measure: Metric::Cosine,
}
);
assert_eq!(after_jq.stages[1].span, Span { start: 15, end: 58 });
assert!(matches!(after_jq.stages[2].value, Stage::Limit(value) if value.get() == 5));
let text = read(r#"notes | similar(.body_vec; "chest pain"; cosine) | limit(5)"#)?;
assert_eq!(
text.stages[0].value,
Stage::Similar {
field: "body_vec".to_owned(),
question: Question::Text("chest pain".to_owned()),
measure: Metric::Cosine,
}
);
let limit = read("notes | limit(5)")?;
assert_eq!(limit.stages.len(), 1);
assert!(matches!(limit.stages[0].value, Stage::Limit(value) if value.get() == 5));
let no_limit = read(r#"notes | similar(.body_vec; record("n-102"); cosine)"#)?;
assert_eq!(no_limit.stages.len(), 1);
assert!(matches!(no_limit.stages[0].value, Stage::Similar { .. }));
let followed_by_jq = read(r#"notes | similar(.body_vec; record("n-102"); cosine) | '.a'"#)?;
assert_eq!(followed_by_jq.stages.len(), 2);
assert_eq!(followed_by_jq.stages[0].value, no_limit.stages[0].value);
assert_eq!(followed_by_jq.stages[1].value, Stage::Jq(".a".to_owned()));
let count = read("notes | count")?;
assert_eq!(count.stages.len(), 1);
assert_eq!(count.stages[0].value, Stage::Count);
Ok(())
}
#[test]
fn a_window_stage_is_refused_by_name() -> Result<(), String> {
let plan = read("notes | window(0; $now)")?;
assert_eq!(plan.stages.len(), 1);
assert_eq!(plan.stages[0].span, super::Span { start: 8, end: 23 });
assert_eq!(
plan.stages[0].value,
Stage::Window(super::window::Window {
from: super::window::Bound::Version(0),
to: super::window::Bound::Now,
})
);
let failure = pipeline("notes | window")
.err()
.ok_or("unclosed window accepted")?;
assert_eq!(failure.code(), "query_window_bound");
assert_eq!(failure.span(), super::Span { start: 8, end: 14 });
assert!(
failure
.to_string()
.contains("window takes from and to in brackets")
);
Ok(())
}
#[test]
fn a_group_by_stage_is_refused_by_name_and_jq_keeps_its_own() -> Result<(), String> {
assert_eq!(
refusal("notes | group_by"),
(
"query_parse",
"`group_by` takes one top-level field in brackets".to_owned()
)
);
let plan = read("notes | group_by(.ward)")?;
assert_eq!(plan.stages.len(), 1);
assert_eq!(
plan.stages[0].value,
Stage::GroupBy {
field: "ward".to_owned(),
}
);
assert_eq!(plan.stages[0].span, super::Span { start: 8, end: 23 });
collection_jq("notes | 'group_by(.ward)'", "notes", "group_by(.ward)")
}
#[test]
fn a_percentile_stage_is_refused_by_name() -> Result<(), String> {
let plan = read("notes | percentile(.admitted; 95)")?;
assert_eq!(plan.stages.len(), 1);
assert_eq!(
plan.stages[0].value,
Stage::Percentile {
field: "admitted".to_owned(),
percentile: serde_json::from_str("95").map_err(|error| error.to_string())?,
}
);
assert_eq!(plan.stages[0].span, super::Span { start: 8, end: 33 });
assert_eq!(
refusal("notes | percentile"),
(
"query_parse",
"`percentile` takes a field and number in brackets".to_owned()
)
);
Ok(())
}
#[test]
fn a_malformed_search_is_a_parse_refusal() {
let cases = [
(
r#"notes | similar(.body_vec; record("n-102"); cosine | limit(5)"#,
"`similar(` is not closed",
),
(
r#"notes | similar(body_vec; record("n-102"); cosine) | limit(5)"#,
"a field is written with a leading dot, as `.body_vec`",
),
(
r#"notes | similar(.body_vec; record("n-102"); cosine) | limit(0)"#,
"`limit(` takes a whole number of 1 or more",
),
(
r#"notes | similar(.body_vec; record("n-102"); cosine) | limit(2.5)"#,
"`limit(` takes a whole number of 1 or more",
),
(
r#"notes | similar(.body_vec; record("n-102"); euclid) | limit(5)"#,
r#""euclid" is not a measure: similar takes cosine, l2 or dot"#,
),
(
r#"notes | similar(.body_vec; record("n-102"; cosine) | limit(5)"#,
"`record(` takes one quoted id and a `)`",
),
];
for (text, words) in cases {
assert_eq!(refusal(text), ("query_parse", words.to_owned()), "{text}");
}
}
#[test]
fn plain_jq_naming_similar_is_still_no_pipeline() {
assert!(matches!(pipeline(".similar"), Err(Refused::Parse(..))));
}