Skip to content

Commit d4033ee

Browse files
Chen-Yuan-LaiCheng-Yuan-LaiIan Laiblaginin
authored
Migrate core test to insta, part1 (#16324)
* feat: migrate core test into insta * feat: enhance tests with snapshot assertions and remove unused macros * feat: rewrite snapshot test in explain_analyze_baseline_metrics * fix: update snapshot test to normalize file paths in explain_analyze_baseline_metrics * refactor: simplify snapshot assertions by removing format calls in optimizer tests * feat: revert the original usage of assert_metrics macro in explain_analyze_baseline_metrics --------- Co-authored-by: Cheng-Yuan-Lai <a186235@g,ail.com> Co-authored-by: Ian Lai <Ian.Lai@senao.com> Co-authored-by: Dmitrii Blaginin <dmitrii@blaginin.me>
1 parent 06631c2 commit d4033ee

File tree

6 files changed

+354
-315
lines changed

6 files changed

+354
-315
lines changed

datafusion/core/tests/expr_api/simplification.rs

Lines changed: 43 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
//! This program demonstrates the DataFusion expression simplification API.
1919
20+
use insta::assert_snapshot;
21+
2022
use arrow::array::types::IntervalDayTime;
2123
use arrow::array::{ArrayRef, Int32Array};
2224
use arrow::datatypes::{DataType, Field, Schema};
@@ -237,11 +239,15 @@ fn to_timestamp_expr_folded() -> Result<()> {
237239
.project(proj)?
238240
.build()?;
239241

240-
let expected = "Projection: TimestampNanosecond(1599566400000000000, None) AS to_timestamp(Utf8(\"2020-09-08T12:00:00+00:00\"))\
241-
\n TableScan: test"
242-
.to_string();
243-
let actual = get_optimized_plan_formatted(plan, &Utc::now());
244-
assert_eq!(expected, actual);
242+
let formatted = get_optimized_plan_formatted(plan, &Utc::now());
243+
let actual = formatted.trim();
244+
assert_snapshot!(
245+
actual,
246+
@r###"
247+
Projection: TimestampNanosecond(1599566400000000000, None) AS to_timestamp(Utf8("2020-09-08T12:00:00+00:00"))
248+
TableScan: test
249+
"###
250+
);
245251
Ok(())
246252
}
247253

@@ -262,11 +268,16 @@ fn now_less_than_timestamp() -> Result<()> {
262268

263269
// Note that constant folder runs and folds the entire
264270
// expression down to a single constant (true)
265-
let expected = "Filter: Boolean(true)\
266-
\n TableScan: test";
267-
let actual = get_optimized_plan_formatted(plan, &time);
268-
269-
assert_eq!(expected, actual);
271+
let formatted = get_optimized_plan_formatted(plan, &time);
272+
let actual = formatted.trim();
273+
274+
assert_snapshot!(
275+
actual,
276+
@r###"
277+
Filter: Boolean(true)
278+
TableScan: test
279+
"###
280+
);
270281
Ok(())
271282
}
272283

@@ -296,11 +307,16 @@ fn select_date_plus_interval() -> Result<()> {
296307

297308
// Note that constant folder runs and folds the entire
298309
// expression down to a single constant (true)
299-
let expected = r#"Projection: Date32("2021-01-09") AS to_timestamp(Utf8("2020-09-08T12:05:00+00:00")) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 0 }")
300-
TableScan: test"#;
301-
let actual = get_optimized_plan_formatted(plan, &time);
302-
303-
assert_eq!(expected, actual);
310+
let formatted = get_optimized_plan_formatted(plan, &time);
311+
let actual = formatted.trim();
312+
313+
assert_snapshot!(
314+
actual,
315+
@r###"
316+
Projection: Date32("2021-01-09") AS to_timestamp(Utf8("2020-09-08T12:05:00+00:00")) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 0 }")
317+
TableScan: test
318+
"###
319+
);
304320
Ok(())
305321
}
306322

@@ -314,10 +330,15 @@ fn simplify_project_scalar_fn() -> Result<()> {
314330

315331
// before simplify: power(t.f, 1.0)
316332
// after simplify: t.f as "power(t.f, 1.0)"
317-
let expected = "Projection: test.f AS power(test.f,Float64(1))\
318-
\n TableScan: test";
319-
let actual = get_optimized_plan_formatted(plan, &Utc::now());
320-
assert_eq!(expected, actual);
333+
let formatter = get_optimized_plan_formatted(plan, &Utc::now());
334+
let actual = formatter.trim();
335+
assert_snapshot!(
336+
actual,
337+
@r###"
338+
Projection: test.f AS power(test.f,Float64(1))
339+
TableScan: test
340+
"###
341+
);
321342
Ok(())
322343
}
323344

@@ -337,9 +358,9 @@ fn simplify_scan_predicate() -> Result<()> {
337358

338359
// before simplify: t.g = power(t.f, 1.0)
339360
// after simplify: t.g = t.f"
340-
let expected = "TableScan: test, full_filters=[g = f]";
341-
let actual = get_optimized_plan_formatted(plan, &Utc::now());
342-
assert_eq!(expected, actual);
361+
let formatted = get_optimized_plan_formatted(plan, &Utc::now());
362+
let actual = formatted.trim();
363+
assert_snapshot!(actual, @"TableScan: test, full_filters=[g = f]");
343364
Ok(())
344365
}
345366

datafusion/core/tests/optimizer/mod.rs

Lines changed: 43 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
//! Tests for the DataFusion SQL query planner that require functions from the
1919
//! datafusion-functions crate.
2020
21+
use insta::assert_snapshot;
2122
use std::any::Any;
2223
use std::collections::HashMap;
2324
use std::sync::Arc;
@@ -56,9 +57,14 @@ fn init() {
5657
#[test]
5758
fn select_arrow_cast() {
5859
let sql = "SELECT arrow_cast(1234, 'Float64') as f64, arrow_cast('foo', 'LargeUtf8') as large";
59-
let expected = "Projection: Float64(1234) AS f64, LargeUtf8(\"foo\") AS large\
60-
\n EmptyRelation";
61-
quick_test(sql, expected);
60+
let plan = test_sql(sql).unwrap();
61+
assert_snapshot!(
62+
plan,
63+
@r#"
64+
Projection: Float64(1234) AS f64, LargeUtf8("foo") AS large
65+
EmptyRelation
66+
"#
67+
);
6268
}
6369
#[test]
6470
fn timestamp_nano_ts_none_predicates() -> Result<()> {
@@ -68,11 +74,15 @@ fn timestamp_nano_ts_none_predicates() -> Result<()> {
6874
// a scan should have the now()... predicate folded to a single
6975
// constant and compared to the column without a cast so it can be
7076
// pushed down / pruned
71-
let expected =
72-
"Projection: test.col_int32\
73-
\n Filter: test.col_ts_nano_none < TimestampNanosecond(1666612093000000000, None)\
74-
\n TableScan: test projection=[col_int32, col_ts_nano_none]";
75-
quick_test(sql, expected);
77+
let plan = test_sql(sql).unwrap();
78+
assert_snapshot!(
79+
plan,
80+
@r"
81+
Projection: test.col_int32
82+
Filter: test.col_ts_nano_none < TimestampNanosecond(1666612093000000000, None)
83+
TableScan: test projection=[col_int32, col_ts_nano_none]
84+
"
85+
);
7686
Ok(())
7787
}
7888

@@ -84,21 +94,30 @@ fn timestamp_nano_ts_utc_predicates() {
8494
// a scan should have the now()... predicate folded to a single
8595
// constant and compared to the column without a cast so it can be
8696
// pushed down / pruned
87-
let expected =
88-
"Projection: test.col_int32\n Filter: test.col_ts_nano_utc < TimestampNanosecond(1666612093000000000, Some(\"+00:00\"))\
89-
\n TableScan: test projection=[col_int32, col_ts_nano_utc]";
90-
quick_test(sql, expected);
97+
let plan = test_sql(sql).unwrap();
98+
assert_snapshot!(
99+
plan,
100+
@r#"
101+
Projection: test.col_int32
102+
Filter: test.col_ts_nano_utc < TimestampNanosecond(1666612093000000000, Some("+00:00"))
103+
TableScan: test projection=[col_int32, col_ts_nano_utc]
104+
"#
105+
);
91106
}
92107

93108
#[test]
94109
fn concat_literals() -> Result<()> {
95110
let sql = "SELECT concat(true, col_int32, false, null, 'hello', col_utf8, 12, 3.4) \
96111
AS col
97112
FROM test";
98-
let expected =
99-
"Projection: concat(Utf8(\"true\"), CAST(test.col_int32 AS Utf8), Utf8(\"falsehello\"), test.col_utf8, Utf8(\"123.4\")) AS col\
100-
\n TableScan: test projection=[col_int32, col_utf8]";
101-
quick_test(sql, expected);
113+
let plan = test_sql(sql).unwrap();
114+
assert_snapshot!(
115+
plan,
116+
@r#"
117+
Projection: concat(Utf8("true"), CAST(test.col_int32 AS Utf8), Utf8("falsehello"), test.col_utf8, Utf8("123.4")) AS col
118+
TableScan: test projection=[col_int32, col_utf8]
119+
"#
120+
);
102121
Ok(())
103122
}
104123

@@ -107,16 +126,15 @@ fn concat_ws_literals() -> Result<()> {
107126
let sql = "SELECT concat_ws('-', true, col_int32, false, null, 'hello', col_utf8, 12, '', 3.4) \
108127
AS col
109128
FROM test";
110-
let expected =
111-
"Projection: concat_ws(Utf8(\"-\"), Utf8(\"true\"), CAST(test.col_int32 AS Utf8), Utf8(\"false-hello\"), test.col_utf8, Utf8(\"12--3.4\")) AS col\
112-
\n TableScan: test projection=[col_int32, col_utf8]";
113-
quick_test(sql, expected);
114-
Ok(())
115-
}
116-
117-
fn quick_test(sql: &str, expected_plan: &str) {
118129
let plan = test_sql(sql).unwrap();
119-
assert_eq!(expected_plan, format!("{plan}"));
130+
assert_snapshot!(
131+
plan,
132+
@r#"
133+
Projection: concat_ws(Utf8("-"), Utf8("true"), CAST(test.col_int32 AS Utf8), Utf8("false-hello"), test.col_utf8, Utf8("12--3.4")) AS col
134+
TableScan: test projection=[col_int32, col_utf8]
135+
"#
136+
);
137+
Ok(())
120138
}
121139

122140
fn test_sql(sql: &str) -> Result<LogicalPlan> {

datafusion/core/tests/sql/aggregates.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use super::*;
1919
use datafusion::common::test_util::batches_to_string;
2020
use datafusion_catalog::MemTable;
2121
use datafusion_common::ScalarValue;
22+
use insta::assert_snapshot;
2223

2324
#[tokio::test]
2425
async fn csv_query_array_agg_distinct() -> Result<()> {

0 commit comments

Comments
 (0)