use super::{Argument, Plan, Question, Source, Span, Stage, pipeline};
use crate::serve::protocol::Metric;
fn read(text: &str) -> Result<Plan, String> {
pipeline(text).map_err(|error| error.to_string())
}
#[test]
fn text_question_is_a_pipeline() -> Result<(), String> {
let plan = read(r#"notes | similar(.v; "chest pain"; cosine) | limit(10)"#)?;
assert_eq!(
plan.source,
Source::Collection {
name: "notes".to_owned(),
arguments: Vec::new()
}
);
assert_eq!(plan.stages.len(), 2);
assert_eq!(
plan.stages[0].value,
Stage::Similar {
field: "v".to_owned(),
question: Question::Text("chest pain".to_owned()),
measure: Metric::Cosine,
}
);
assert!(matches!(plan.stages[1].value, Stage::Limit(value) if value.get() == 10));
Ok(())
}
#[test]
fn collection_arguments_sort_and_limit_are_a_pipeline() -> Result<(), String> {
let plan =
read(r#"patients(ward: "4B" and status: "open") | sort(.admitted; desc) | limit(20)"#)?;
assert_eq!(
plan.source,
Source::Collection {
name: "patients".to_owned(),
arguments: vec![
Argument {
field: "ward".to_owned(),
value: serde_json::json!("4B")
},
Argument {
field: "status".to_owned(),
value: serde_json::json!("open")
},
],
}
);
assert_eq!(plan.stages.len(), 2);
assert_eq!(
plan.stages[0].value,
Stage::Sort {
field: "admitted".to_owned(),
descending: true
}
);
assert!(matches!(plan.stages[1].value, Stage::Limit(value) if value.get() == 20));
Ok(())
}
#[test]
fn limit_can_follow_jq() -> Result<(), String> {
let plan = read("patients | '.name' | limit(5)")?;
assert_eq!(plan.stages.len(), 2);
assert_eq!(plan.stages[0].value, Stage::Jq(".name".to_owned()));
assert!(matches!(plan.stages[1].value, Stage::Limit(value) if value.get() == 5));
Ok(())
}
#[test]
fn count_is_a_pipeline() -> Result<(), String> {
let plan = read("patients | count")?;
assert_eq!(plan.stages.len(), 1);
assert_eq!(plan.stages[0].value, Stage::Count);
Ok(())
}
#[test]
fn nothing_follows_count() {
let refusal = pipeline("patients | count | limit(1)").err();
assert!(refusal.is_some_and(|error| {
error.to_string().contains("nothing follows count")
&& error.span() == Span { start: 19, end: 27 }
}));
}
#[test]
fn similar_after_jq_keeps_its_arguments_and_position() -> Result<(), String> {
let text = "patients | '.a' | similar(.v; record(\"p0\"); dot)";
let plan = read(text)?;
assert_eq!(plan.stages[0].value, Stage::Jq(".a".to_owned()));
assert_eq!(
plan.stages[1].value,
Stage::Similar {
field: "v".to_owned(),
question: Question::Record("p0".to_owned()),
measure: Metric::Dot,
}
);
assert_eq!(
plan.stages[1].span,
Span {
start: 18,
end: text.len()
}
);
Ok(())
}
#[test]
fn sort_by_stays_not_yet_with_its_position() {
let refusal = pipeline("patients | sort_by(.x)").err();
assert!(refusal.is_some_and(|error| {
error.code() == "query_stage_not_available"
&& error.to_string().contains("sort_by")
&& error.to_string().contains("bytes 11..18")
&& error.span() == Span { start: 11, end: 18 }
}));
}
#[test]
fn stages_record_their_byte_spans() -> Result<(), String> {
let text = concat!(
" \u{2003}patients | ",
r#"similar(.v; "café"; cosine)"#,
" | limit(10) | sort(.admitted; desc) | ",
r#"'select(.name == "O''Brien")'"#,
" | limit(5) | count \n"
);
let plan = read(text)?;
let written = [
r#"similar(.v; "café"; cosine)"#,
"limit(10)",
"sort(.admitted; desc)",
r#"'select(.name == "O''Brien")'"#,
"limit(5)",
"count",
];
assert_eq!(plan.stages.len(), written.len());
for (stage, written) in plan.stages.iter().zip(written) {
let start = text.find(written).ok_or("missing stage")?;
assert_eq!(
stage.span,
Span {
start,
end: start + written.len()
}
);
assert_eq!(text.get(stage.span.start..stage.span.end), Some(written));
}
let original = read(" patients | '.name' ")?;
assert_eq!(original.stages[0].span, Span { start: 13, end: 20 });
let all = read(" 'é' | limit(1)")?;
assert_eq!(all.source, Source::All);
assert_eq!(all.stages[0].span, Span { start: 1, end: 5 });
assert_eq!(all.stages[1].span, Span { start: 8, end: 16 });
Ok(())
}
#[test]
fn served_stages_keep_their_words_and_complete_byte_spans() -> Result<(), String> {
for (word, query) in [
("rg", r#"patients | rg("4B")"#),
("match", "patients | match (c:patients)"),
("window", "patients | window(0; $now)"),
(
"rerank",
r#"patients | rerank(.embedding; record("p0"); dot)"#,
),
("{", "patients | { name }"),
("group_by", "patients | group_by(.ward)"),
("percentile", "patients | percentile(.admitted; 95)"),
] {
let plan = read(query)?;
assert_eq!(plan.stages.len(), 1, "{query}");
let stage = plan.stages.first().ok_or("served stage missing")?;
let start = query.find(word).ok_or("served word missing")?;
assert_eq!(
stage.span,
Span {
start,
end: query.len()
}
);
assert!(
matches!(
(&stage.value, word),
(Stage::Rg { .. }, "rg")
| (Stage::Match { .. }, "match")
| (Stage::Window(_), "window")
| (Stage::Rerank { .. }, "rerank")
| (Stage::Shape(_), "{")
| (Stage::GroupBy { .. }, "group_by")
| (Stage::Percentile { .. }, "percentile")
),
"{query}"
);
}
Ok(())
}
#[test]
fn similar_after_jq_keeps_malformed_argument_refusal_and_position() -> Result<(), String> {
let text = "patients | '.a' | similar(...)";
let refusal = pipeline(text).err().ok_or("malformed similar parsed")?;
assert_eq!(refusal.code(), "query_parse");
assert_eq!(refusal.span(), Span { start: 27, end: 29 });
assert_eq!(
text.get(refusal.span().start..refusal.span().end),
Some("..")
);
assert!(
refusal
.to_string()
.contains("a field has a name after its dot")
);
Ok(())
}