arrow/util/
test_util.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Utils to make testing easier
19
20use rand::{RngExt, SeedableRng, rngs::StdRng};
21use std::{env, error::Error, fs, io::Write, path::PathBuf};
22
23/// Returns a vector of size `n`, filled with randomly generated bytes.
24pub fn random_bytes(n: usize) -> Vec<u8> {
25    let mut result = vec![];
26    let mut rng = seedable_rng();
27    for _ in 0..n {
28        result.push(rng.random_range(0..255));
29    }
30    result
31}
32
33/// Returns fixed seedable RNG
34pub fn seedable_rng() -> StdRng {
35    StdRng::seed_from_u64(42)
36}
37
38/// Returns file handle for a temp file in 'target' directory with a provided content
39///
40/// TODO: Originates from `parquet` utils, can be merged in [ARROW-4064]
41pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
42    // build tmp path to a file in "target/debug/testdata"
43    let dir = env::current_dir().unwrap().join("target/debug/testdata");
44    fs::create_dir_all(&dir).unwrap();
45    let path_buf = dir.join(file_name);
46
47    // write file content
48    let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
49    tmp_file.write_all(content).unwrap();
50    tmp_file.sync_all().unwrap();
51
52    // return file handle for both read and write
53    let file = fs::OpenOptions::new()
54        .read(true)
55        .write(true)
56        .open(path_buf.as_path());
57    assert!(file.is_ok());
58    file.unwrap()
59}
60
61/// Returns the arrow test data directory, which is by default stored
62/// in a git submodule rooted at `arrow/testing/data`.
63///
64/// The default can be overridden by the optional environment
65/// variable `ARROW_TEST_DATA`
66///
67/// panics when the directory can not be found.
68///
69/// Example:
70/// ```
71/// let testdata = arrow::util::test_util::arrow_test_data();
72/// let csvdata = format!("{}/csv/aggregate_test_100.csv", testdata);
73/// assert!(std::path::PathBuf::from(csvdata).exists());
74/// ```
75pub fn arrow_test_data() -> String {
76    match get_data_dir("ARROW_TEST_DATA", "../testing/data") {
77        Ok(pb) => pb.display().to_string(),
78        Err(err) => panic!("failed to get arrow data dir: {err}"),
79    }
80}
81
82/// Returns the parquest test data directory, which is by default
83/// stored in a git submodule rooted at
84/// `arrow/parquest-testing/data`.
85///
86/// The default can be overridden by the optional environment variable
87/// `PARQUET_TEST_DATA`
88///
89/// panics when the directory can not be found.
90///
91/// Example:
92/// ```
93/// let testdata = arrow::util::test_util::parquet_test_data();
94/// let filename = format!("{}/binary.parquet", testdata);
95/// assert!(std::path::PathBuf::from(filename).exists());
96/// ```
97pub fn parquet_test_data() -> String {
98    match get_data_dir("PARQUET_TEST_DATA", "../parquet-testing/data") {
99        Ok(pb) => pb.display().to_string(),
100        Err(err) => panic!("failed to get parquet data dir: {err}"),
101    }
102}
103
104/// Returns a directory path for finding test data.
105///
106/// udf_env: name of an environment variable
107///
108/// submodule_dir: fallback path (relative to CARGO_MANIFEST_DIR)
109///
110///  Returns either:
111/// The path referred to in `udf_env` if that variable is set and refers to a directory
112/// The submodule_data directory relative to CARGO_MANIFEST_PATH
113fn get_data_dir(udf_env: &str, submodule_data: &str) -> Result<PathBuf, Box<dyn Error>> {
114    // Try user defined env.
115    if let Ok(dir) = env::var(udf_env) {
116        let trimmed = dir.trim().to_string();
117        if !trimmed.is_empty() {
118            let pb = PathBuf::from(trimmed);
119            if pb.is_dir() {
120                return Ok(pb);
121            } else {
122                return Err(format!(
123                    "the data dir `{}` defined by env {} not found",
124                    pb.display(),
125                    udf_env
126                )
127                .into());
128            }
129        }
130    }
131
132    // The env is undefined or its value is trimmed to empty, let's try default dir.
133
134    // env "CARGO_MANIFEST_DIR" is "the directory containing the manifest of your package",
135    // set by `cargo run` or `cargo test`, see:
136    // https://doc.rust-lang.org/cargo/reference/environment-variables.html
137    let dir = env!("CARGO_MANIFEST_DIR");
138
139    let pb = PathBuf::from(dir).join(submodule_data);
140    if pb.is_dir() {
141        Ok(pb)
142    } else {
143        Err(format!(
144            "env `{}` is undefined or has empty value, and the pre-defined data dir `{}` not found\n\
145             HINT: try running `git submodule update --init`",
146            udf_env,
147            pb.display(),
148        ).into())
149    }
150}
151
152/// An iterator that is untruthful about its actual length
153#[derive(Debug, Clone)]
154pub struct BadIterator<T> {
155    /// where the iterator currently is
156    cur: usize,
157    /// How many items will this iterator *actually* make
158    limit: usize,
159    /// How many items this iterator claims it will make
160    claimed: usize,
161    /// The items to return. If there are fewer items than `limit`
162    /// they will be repeated
163    pub items: Vec<T>,
164}
165
166impl<T> BadIterator<T> {
167    /// Create a new iterator for `<limit>` items, but that reports to
168    /// produce `<claimed>` items. Must provide at least 1 item.
169    pub fn new(limit: usize, claimed: usize, items: Vec<T>) -> Self {
170        assert!(!items.is_empty());
171        Self {
172            cur: 0,
173            limit,
174            claimed,
175            items,
176        }
177    }
178}
179
180impl<T: Clone> Iterator for BadIterator<T> {
181    type Item = T;
182
183    fn next(&mut self) -> Option<Self::Item> {
184        if self.cur < self.limit {
185            let next_item_idx = self.cur % self.items.len();
186            let next_item = self.items[next_item_idx].clone();
187            self.cur += 1;
188            Some(next_item)
189        } else {
190            None
191        }
192    }
193
194    /// report whatever the iterator says to
195    fn size_hint(&self) -> (usize, Option<usize>) {
196        (0, Some(self.claimed))
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    #[test]
205    fn test_data_dir() {
206        let udf_env = "get_data_dir";
207        let cwd = env::current_dir().unwrap();
208
209        let existing_pb = cwd.join("..");
210        let existing = existing_pb.display().to_string();
211        let existing_str = existing.as_str();
212
213        let non_existing = cwd.join("non-existing-dir").display().to_string();
214        let non_existing_str = non_existing.as_str();
215
216        unsafe { env::set_var(udf_env, non_existing_str) };
217        let res = get_data_dir(udf_env, existing_str);
218        assert!(res.is_err());
219
220        unsafe { env::set_var(udf_env, "") };
221        let res = get_data_dir(udf_env, existing_str);
222        assert!(res.is_ok());
223        assert_eq!(res.unwrap(), existing_pb);
224
225        unsafe { env::set_var(udf_env, " ") };
226        let res = get_data_dir(udf_env, existing_str);
227        assert!(res.is_ok());
228        assert_eq!(res.unwrap(), existing_pb);
229
230        unsafe { env::set_var(udf_env, existing_str) };
231        let res = get_data_dir(udf_env, existing_str);
232        assert!(res.is_ok());
233        assert_eq!(res.unwrap(), existing_pb);
234
235        unsafe { env::remove_var(udf_env) };
236        let res = get_data_dir(udf_env, non_existing_str);
237        assert!(res.is_err());
238
239        let res = get_data_dir(udf_env, existing_str);
240        assert!(res.is_ok());
241        assert_eq!(res.unwrap(), existing_pb);
242    }
243
244    #[test]
245    fn test_happy() {
246        let res = arrow_test_data();
247        assert!(PathBuf::from(res).is_dir());
248
249        let res = parquet_test_data();
250        assert!(PathBuf::from(res).is_dir());
251    }
252}
arrow/util/test_util.rs

arrow/util/
test_util.rs
