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/// Example:
68/// ```
69/// let testdata = arrow::util::test_util::arrow_test_data();
70/// let csvdata = format!("{}/csv/aggregate_test_100.csv", testdata);
71/// assert!(std::path::PathBuf::from(csvdata).exists());
72/// ```
73///
74/// # Panics
75///
76/// Panics if the directory can not be found
77pub fn arrow_test_data() -> String {
78    match get_data_dir("ARROW_TEST_DATA", "../testing/data") {
79        Ok(pb) => pb.display().to_string(),
80        Err(err) => panic!("failed to get arrow data dir: {err}"),
81    }
82}
83
84/// Returns the parquest test data directory, which is by default
85/// stored in a git submodule rooted at
86/// `arrow/parquest-testing/data`.
87///
88/// The default can be overridden by the optional environment variable
89/// `PARQUET_TEST_DATA`
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/// ```
97///
98/// # Panics
99///
100/// Panics if the directory can not be found
101pub fn parquet_test_data() -> String {
102    match get_data_dir("PARQUET_TEST_DATA", "../parquet-testing/data") {
103        Ok(pb) => pb.display().to_string(),
104        Err(err) => panic!("failed to get parquet data dir: {err}"),
105    }
106}
107
108/// Returns a directory path for finding test data.
109///
110/// udf_env: name of an environment variable
111///
112/// submodule_dir: fallback path (relative to CARGO_MANIFEST_DIR)
113///
114///  Returns either:
115/// The path referred to in `udf_env` if that variable is set and refers to a directory
116/// The submodule_data directory relative to CARGO_MANIFEST_PATH
117fn get_data_dir(udf_env: &str, submodule_data: &str) -> Result<PathBuf, Box<dyn Error>> {
118    // Try user defined env.
119    if let Ok(dir) = env::var(udf_env) {
120        let trimmed = dir.trim().to_string();
121        if !trimmed.is_empty() {
122            let pb = PathBuf::from(trimmed);
123            return if pb.is_dir() {
124                Ok(pb)
125            } else {
126                Err(format!(
127                    "the data dir `{}` defined by env {} not found",
128                    pb.display(),
129                    udf_env
130                )
131                .into())
132            };
133        }
134    }
135
136    // The env is undefined or its value is trimmed to empty, let's try default dir.
137
138    // env "CARGO_MANIFEST_DIR" is "the directory containing the manifest of your package",
139    // set by `cargo run` or `cargo test`, see:
140    // https://doc.rust-lang.org/cargo/reference/environment-variables.html
141    let dir = env!("CARGO_MANIFEST_DIR");
142
143    let pb = PathBuf::from(dir).join(submodule_data);
144    if pb.is_dir() {
145        Ok(pb)
146    } else {
147        Err(format!(
148            "env `{}` is undefined or has empty value, and the pre-defined data dir `{}` not found\n\
149             HINT: try running `git submodule update --init`",
150            udf_env,
151            pb.display(),
152        ).into())
153    }
154}
155
156/// An iterator that is untruthful about its actual length
157#[derive(Debug, Clone)]
158pub struct BadIterator<T> {
159    /// where the iterator currently is
160    cur: usize,
161    /// How many items will this iterator *actually* make
162    limit: usize,
163    /// How many items this iterator claims it will make
164    claimed: usize,
165    /// The items to return. If there are fewer items than `limit`
166    /// they will be repeated
167    pub items: Vec<T>,
168}
169
170impl<T> BadIterator<T> {
171    /// Create a new iterator for `<limit>` items, but that reports to
172    /// produce `<claimed>` items. Must provide at least 1 item.
173    ///
174    /// # Panics
175    ///
176    /// Panics if `items` is empty
177    pub fn new(limit: usize, claimed: usize, items: Vec<T>) -> Self {
178        assert!(!items.is_empty());
179        Self {
180            cur: 0,
181            limit,
182            claimed,
183            items,
184        }
185    }
186}
187
188impl<T: Clone> Iterator for BadIterator<T> {
189    type Item = T;
190
191    fn next(&mut self) -> Option<Self::Item> {
192        if self.cur < self.limit {
193            let next_item_idx = self.cur % self.items.len();
194            let next_item = self.items[next_item_idx].clone();
195            self.cur += 1;
196            Some(next_item)
197        } else {
198            None
199        }
200    }
201
202    /// report whatever the iterator says to
203    fn size_hint(&self) -> (usize, Option<usize>) {
204        (0, Some(self.claimed))
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn test_data_dir() {
214        let udf_env = "get_data_dir";
215        let cwd = env::current_dir().unwrap();
216
217        let existing_pb = cwd.join("..");
218        let existing = existing_pb.display().to_string();
219        let existing_str = existing.as_str();
220
221        let non_existing = cwd.join("non-existing-dir").display().to_string();
222        let non_existing_str = non_existing.as_str();
223
224        unsafe { env::set_var(udf_env, non_existing_str) };
225        let res = get_data_dir(udf_env, existing_str);
226        assert!(res.is_err());
227
228        unsafe { env::set_var(udf_env, "") };
229        let res = get_data_dir(udf_env, existing_str);
230        assert!(res.is_ok());
231        assert_eq!(res.unwrap(), existing_pb);
232
233        unsafe { env::set_var(udf_env, " ") };
234        let res = get_data_dir(udf_env, existing_str);
235        assert!(res.is_ok());
236        assert_eq!(res.unwrap(), existing_pb);
237
238        unsafe { env::set_var(udf_env, existing_str) };
239        let res = get_data_dir(udf_env, existing_str);
240        assert!(res.is_ok());
241        assert_eq!(res.unwrap(), existing_pb);
242
243        unsafe { env::remove_var(udf_env) };
244        let res = get_data_dir(udf_env, non_existing_str);
245        assert!(res.is_err());
246
247        let res = get_data_dir(udf_env, existing_str);
248        assert!(res.is_ok());
249        assert_eq!(res.unwrap(), existing_pb);
250    }
251
252    #[test]
253    fn test_happy() {
254        let res = arrow_test_data();
255        assert!(PathBuf::from(res).is_dir());
256
257        let res = parquet_test_data();
258        assert!(PathBuf::from(res).is_dir());
259    }
260}
arrow/util/test_util.rs

arrow/util/
test_util.rs
