Skip to content

Commit a3980db

Browse files
blagininalamb
andauthored
Use insta for DataFrame tests (#15165)
* Add insta for df testing * Do not use `get_plan_string` * Switch from `assert_batches_eq` * Switch from `assert_batches_sorted_eq` * Switch from `assert_eq` * Format toml --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent e37d580 commit a3980db

File tree

6 files changed

+1824
-1642
lines changed

6 files changed

+1824
-1642
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ parquet = { version = "54.2.1", default-features = false, features = [
155155
pbjson = { version = "0.7.0" }
156156
pbjson-types = "0.7"
157157
# Should match arrow-flight's version of prost.
158+
insta = { version = "1.41.1", features = ["glob", "filters"] }
158159
prost = "0.13.1"
159160
rand = "0.8.5"
160161
recursive = "0.1.1"

datafusion-cli/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ url = { workspace = true }
6767
[dev-dependencies]
6868
assert_cmd = "2.0"
6969
ctor = { workspace = true }
70-
insta = { version = "1.41.1", features = ["glob", "filters"] }
70+
insta = { workspace = true }
7171
insta-cmd = "0.6.0"
7272
predicates = "3.0"
7373
rstest = { workspace = true }

datafusion/common/src/test_util.rs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@
1717

1818
//! Utility functions to make testing DataFusion based crates easier
1919
20+
use crate::arrow::util::pretty::pretty_format_batches_with_options;
21+
use crate::format::DEFAULT_FORMAT_OPTIONS;
22+
use arrow::array::RecordBatch;
2023
use std::{error::Error, path::PathBuf};
2124

2225
/// Compares formatted output of a record batch with an expected
@@ -73,6 +76,31 @@ macro_rules! assert_batches_eq {
7376
};
7477
}
7578

79+
pub fn batches_to_string(batches: &[RecordBatch]) -> String {
80+
let actual = pretty_format_batches_with_options(batches, &DEFAULT_FORMAT_OPTIONS)
81+
.unwrap()
82+
.to_string();
83+
84+
actual.trim().to_string()
85+
}
86+
87+
pub fn batches_to_sort_string(batches: &[RecordBatch]) -> String {
88+
let actual_lines =
89+
pretty_format_batches_with_options(batches, &DEFAULT_FORMAT_OPTIONS)
90+
.unwrap()
91+
.to_string();
92+
93+
let mut actual_lines: Vec<&str> = actual_lines.trim().lines().collect();
94+
95+
// sort except for header + footer
96+
let num_lines = actual_lines.len();
97+
if num_lines > 3 {
98+
actual_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
99+
}
100+
101+
actual_lines.join("\n")
102+
}
103+
76104
/// Compares formatted output of a record batch with an expected
77105
/// vector of strings in a way that order does not matter.
78106
/// This is a macro so errors appear on the correct line

datafusion/core/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@ datafusion-functions-window-common = { workspace = true }
152152
datafusion-physical-optimizer = { workspace = true }
153153
doc-comment = { workspace = true }
154154
env_logger = { workspace = true }
155+
insta = { workspace = true }
155156
paste = "^1.0"
156157
rand = { workspace = true, features = ["small_rng"] }
157158
rand_distr = "0.4.3"

0 commit comments

Comments
 (0)