Skip to content

Commit c386f20

Browse files
joseph-isaacsjoseph-isaacs
committed
fix
Signed-off-by: Joe Isaacs <[email protected]>
1 parent 88ca847 commit c386f20

File tree

3 files changed

+102
-43
lines changed

3 files changed

+102
-43
lines changed

‎vortex-layout/src/display.rs‎

Lines changed: 43 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ mod tests {
229229
use vortex_dtype::PType;
230230
use vortex_dtype::StructFields;
231231
use vortex_io::runtime::single::block_on;
232+
use vortex_utils::env::EnvVarGuard;
232233

233234
use crate::IntoLayout;
234235
use crate::OwnedLayoutChildren;
@@ -244,6 +245,7 @@ mod tests {
244245
/// Test display_tree with inline array_tree metadata (no segment source needed).
245246
#[test]
246247
fn test_display_tree_inline_array_tree() {
248+
let _guard = EnvVarGuard::set("FLAT_LAYOUT_INLINE_ARRAY_NODE", "1");
247249
block_on(|handle| async move {
248250
let ctx = ArrayContext::empty();
249251
let segments = Arc::new(TestSegments::default());
@@ -333,6 +335,8 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5
333335
/// Test display_tree_with_segments using async segment source to fetch buffer sizes.
334336
#[test]
335337
fn test_display_tree_with_segment_source() {
338+
// Ensure inline array node is disabled for this test
339+
let _guard = EnvVarGuard::remove("FLAT_LAYOUT_INLINE_ARRAY_NODE");
336340
block_on(|handle| async move {
337341
let ctx = ArrayContext::empty();
338342
let segments = Arc::new(TestSegments::default());
@@ -380,31 +384,26 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5
380384

381385
let expected = "\
382386
vortex.chunked, dtype: i32, children: 2, rows: 10
383-
├── [0]: vortex.flat, dtype: i32, metadata: 98 bytes, rows: 5, segment 0, buffers=[20B], total=20B
384-
└── [1]: vortex.flat, dtype: i32, metadata: 98 bytes, rows: 5, segment 1, buffers=[20B], total=20B
387+
├── [0]: vortex.flat, dtype: i32, rows: 5, segment 0, buffers=[20B], total=20B
388+
└── [1]: vortex.flat, dtype: i32, rows: 5, segment 1, buffers=[20B], total=20B
385389
";
386-
assert_eq!(format!("{}", output), expected);
390+
assert_eq!(output.to_string(), expected);
387391
})
388392
}
389393

390394
/// Test display_array_tree with inline array node metadata.
391395
#[test]
392396
fn test_display_array_tree_with_inline_node() {
393-
// Set the env var to enable inlined array node
394-
// SAFETY: This test is single-threaded and we're only setting an env var
395-
unsafe {
396-
std::env::set_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", "1");
397-
}
398-
399-
block_on(|handle| async {
400-
let ctx = ArrayContext::empty();
401-
let segments = Arc::new(TestSegments::default());
402-
let (ptr, eof) = SequenceId::root().split();
397+
let _guard = EnvVarGuard::set("FLAT_LAYOUT_INLINE_ARRAY_NODE", "1");
403398

404-
// Create a simple primitive array
405-
let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid);
399+
let ctx = ArrayContext::empty();
400+
let segments = Arc::new(TestSegments::default());
401+
let (ptr, eof) = SequenceId::root().split();
406402

407-
let layout = FlatLayoutStrategy::default()
403+
// Create a simple primitive array
404+
let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid);
405+
let layout = block_on(|handle| async {
406+
FlatLayoutStrategy::default()
408407
.write_stream(
409408
ctx.clone(),
410409
segments.clone(),
@@ -413,39 +412,40 @@ vortex.chunked, dtype: i32, children: 2, rows: 10
413412
handle,
414413
)
415414
.await
416-
.unwrap();
415+
.unwrap()
416+
});
417417

418-
let flat_layout = layout.as_::<FlatVTable>();
418+
let flat_layout = layout.as_::<FlatVTable>();
419419

420-
let array_tree = flat_layout
421-
.array_tree()
422-
.expect("array_tree should be populated when FLAT_LAYOUT_INLINE_ARRAY_NODE is set");
420+
let array_tree = flat_layout
421+
.array_tree()
422+
.expect("array_tree should be populated when FLAT_LAYOUT_INLINE_ARRAY_NODE is set");
423423

424-
let parts = ArrayParts::from_array_tree(array_tree.as_ref().to_vec())
425-
.expect("should parse array_tree");
426-
assert_eq!(parts.buffer_lengths(), vec![20]); // 5 i32 values = 20 bytes
424+
let parts = ArrayParts::from_array_tree(array_tree.as_ref().to_vec())
425+
.expect("should parse array_tree");
426+
assert_eq!(parts.buffer_lengths(), vec![20]); // 5 i32 values = 20 bytes
427427

428-
assert_eq!(
429-
layout.display_tree().to_string(),
430-
"\
428+
assert_eq!(
429+
layout.display_tree().to_string(),
430+
"\
431431
vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B
432432
"
433-
);
434-
})
433+
);
435434
}
436435

437436
/// Test display_tree without inline array node (shows segment ID).
438437
#[test]
439438
fn test_display_tree_without_inline_node() {
440-
block_on(|handle| async {
441-
let ctx = ArrayContext::empty();
442-
let segments = Arc::new(TestSegments::default());
443-
let (ptr, eof) = SequenceId::root().split();
439+
let _guard = EnvVarGuard::set("FLAT_LAYOUT_INLINE_ARRAY_NODE", "1");
444440

445-
// Create a simple primitive array
446-
let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable);
441+
let ctx = ArrayContext::empty();
442+
let segments = Arc::new(TestSegments::default());
443+
let (ptr, eof) = SequenceId::root().split();
447444

448-
let layout = FlatLayoutStrategy::default()
445+
// Create a simple primitive array
446+
let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable);
447+
let layout = block_on(|handle| async {
448+
FlatLayoutStrategy::default()
449449
.write_stream(
450450
ctx,
451451
segments.clone(),
@@ -454,15 +454,15 @@ vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B
454454
handle,
455455
)
456456
.await
457-
.unwrap();
457+
.unwrap()
458+
});
458459

459-
// Test display_tree exact output (with inline array_tree enabled by env var from other test)
460-
assert_eq!(
461-
layout.display_tree().to_string(),
462-
"\
460+
// Test display_tree exact output (with inline array_tree enabled by env var from other test)
461+
assert_eq!(
462+
layout.display_tree().to_string(),
463+
"\
463464
vortex.flat, dtype: i64, segment 0, buffers=[24B], total=24B
464465
"
465-
);
466-
})
466+
);
467467
}
468468
}

‎vortex-utils/src/env.rs‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Environment variable utilities for testing.
5+
6+
/// RAII guard to set/remove an environment variable for the duration of a scope.
7+
///
8+
/// Removes the variable when dropped, ensuring test isolation.
9+
///
10+
/// # Example
11+
///
12+
/// ```
13+
/// use vortex_utils::env::EnvVarGuard;
14+
///
15+
/// // Set an env var for the duration of this scope
16+
/// let _guard = EnvVarGuard::set("MY_TEST_VAR", "1");
17+
/// assert_eq!(std::env::var("MY_TEST_VAR").ok(), Some("1".to_string()));
18+
///
19+
/// // Or remove an env var
20+
/// let _guard2 = EnvVarGuard::remove("MY_TEST_VAR");
21+
/// assert!(std::env::var("MY_TEST_VAR").is_err());
22+
/// ```
23+
///
24+
/// # Safety
25+
///
26+
/// Environment variable modification is inherently unsafe in multi-threaded contexts.
27+
/// This guard is intended for use in tests that are run serially or where env var
28+
/// races are acceptable.
29+
pub struct EnvVarGuard(&'static str);
30+
31+
impl EnvVarGuard {
32+
/// Set an environment variable for the duration of this guard's lifetime.
33+
pub fn set(key: &'static str, value: &str) -> Self {
34+
// SAFETY: Tests are run serially or we accept env var race risk.
35+
unsafe {
36+
std::env::set_var(key, value);
37+
}
38+
Self(key)
39+
}
40+
41+
/// Remove an environment variable for the duration of this guard's lifetime.
42+
pub fn remove(key: &'static str) -> Self {
43+
// SAFETY: Tests are run serially or we accept env var race risk.
44+
unsafe {
45+
std::env::remove_var(key);
46+
}
47+
Self(key)
48+
}
49+
}
50+
51+
impl Drop for EnvVarGuard {
52+
fn drop(&mut self) {
53+
// SAFETY: See above.
54+
unsafe {
55+
std::env::remove_var(self.0);
56+
}
57+
}
58+
}

‎vortex-utils/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,4 @@ pub mod aliases;
99
pub mod debug_with;
1010
#[cfg(feature = "dyn-traits")]
1111
pub mod dyn_traits;
12+
pub mod env;

0 commit comments

Comments
 (0)