Respect ARROW_TEST_DATA in apache-avro tests (#4950)

apache · Oct 18, 2023 · 51ac6fe · 51ac6fe
1 parent 6e332b8
commit 51ac6fe
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/arrow-avro/src/lib.rs b/arrow-avro/src/lib.rs
@@ -26,3 +26,13 @@ pub mod reader;
 mod schema;
 
 mod compression;
+
+#[cfg(test)]
+mod test_util {
+    pub fn arrow_test_data(path: &str) -> String {
+        match std::env::var("ARROW_TEST_DATA") {
+            Ok(dir) => format!("{dir}/{path}"),
+            Err(_) => format!("../testing/data/{path}"),
+        }
+    }
+}
diff --git a/arrow-avro/src/reader/header.rs b/arrow-avro/src/reader/header.rs
@@ -240,6 +240,7 @@ mod test {
     use super::*;
     use crate::reader::read_header;
     use crate::schema::SCHEMA_METADATA_KEY;
+    use crate::test_util::arrow_test_data;
     use std::fs::File;
     use std::io::{BufRead, BufReader};
 
@@ -266,7 +267,7 @@ mod test {
 
     #[test]
     fn test_header() {
-        let header = decode_file("../testing/data/avro/alltypes_plain.avro");
+        let header = decode_file(&arrow_test_data("avro/alltypes_plain.avro"));
         let schema_json = header.get(SCHEMA_METADATA_KEY).unwrap();
         let expected = br#"{"type":"record","name":"topLevelRecord","fields":[{"name":"id","type":["int","null"]},{"name":"bool_col","type":["boolean","null"]},{"name":"tinyint_col","type":["int","null"]},{"name":"smallint_col","type":["int","null"]},{"name":"int_col","type":["int","null"]},{"name":"bigint_col","type":["long","null"]},{"name":"float_col","type":["float","null"]},{"name":"double_col","type":["double","null"]},{"name":"date_string_col","type":["bytes","null"]},{"name":"string_col","type":["bytes","null"]},{"name":"timestamp_col","type":[{"type":"long","logicalType":"timestamp-micros"},"null"]}]}"#;
         assert_eq!(schema_json, expected);
@@ -276,7 +277,7 @@ mod test {
             226966037233754408753420635932530907102
         );
 
-        let header = decode_file("../testing/data/avro/fixed_length_decimal.avro");
+        let header = decode_file(&arrow_test_data("avro/fixed_length_decimal.avro"));
         let schema_json = header.get(SCHEMA_METADATA_KEY).unwrap();
         let expected = br#"{"type":"record","name":"topLevelRecord","fields":[{"name":"value","type":[{"type":"fixed","name":"fixed","namespace":"topLevelRecord.value","size":11,"logicalType":"decimal","precision":25,"scale":2},"null"]}]}"#;
         assert_eq!(schema_json, expected);

diff --git a/arrow-avro/src/reader/mod.rs b/arrow-avro/src/reader/mod.rs
@@ -76,12 +76,13 @@ fn read_blocks<R: BufRead>(
 #[cfg(test)]
 mod test {
     use crate::reader::{read_blocks, read_header};
+    use crate::test_util::arrow_test_data;
     use std::fs::File;
     use std::io::BufReader;
 
     #[test]
     fn test_mux() {
-        let file = File::open("../testing/data/avro/alltypes_plain.avro").unwrap();
+        let file = File::open(arrow_test_data("avro/alltypes_plain.avro")).unwrap();
         let mut reader = BufReader::new(file);
         let header = read_header(&mut reader).unwrap();
         for result in read_blocks(reader) {