Move all test utils to src/ and hide them behind an "integration" feature. This way, other crates and integration tests as well can use it

Author: gabotechs

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ jobs:
         with:
           lfs: true
       - uses: ./.github/actions/setup
-      - run: cargo test
+      - run: cargo test --features integration
 
   format-check:
     runs-on: ubuntu-latest

Cargo.toml

@@ -23,5 +23,13 @@ prost = "0.13.5"
 rand = "0.8.5"
 object_store = "0.12.3"
 
+# integration_tests deps
+insta = { version = "1.43.1", features = ["filters"], optional = true }
+
+[features]
+integration = [
+    "insta"
+]
+
 [dev-dependencies]
 insta = { version = "1.43.1", features = ["filters"] }

src/flight_service/stream_partitioner_registry.rs

@@ -35,7 +35,7 @@ impl StreamPartitionerRegistry {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::test_utils::MockExec;
+    use crate::test_utils::mock_exec::MockExec;
     use datafusion::arrow::array::{RecordBatch, UInt32Array};
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::execution::TaskContext;

src/lib.rs

@@ -3,16 +3,18 @@ mod common;
 mod composed_extension_codec;
 mod errors;
 mod flight_service;
+mod physical_optimizer;
 mod plan;
-#[cfg(test)]
-mod test_utils;
-
-pub mod physical_optimizer;
-pub mod stage;
-pub mod task;
+mod stage;
+mod task;
 mod user_provided_codec;
 
+#[cfg(any(feature = "integration", test))]
+pub mod test_utils;
+
 pub use channel_manager::{BoxCloneSyncChannel, ChannelManager, ChannelResolver};
 pub use flight_service::{ArrowFlightEndpoint, SessionBuilder};
+pub use physical_optimizer::DistributedPhysicalOptimizerRule;
 pub use plan::ArrowFlightReadExec;
+pub use stage::{display_stage_graphviz, ExecutionStage};
 pub use user_provided_codec::{add_user_codec, with_user_codec};

src/physical_optimizer.rs

@@ -153,7 +153,7 @@ impl DistributedPhysicalOptimizerRule {
 mod tests {
     use crate::assert_snapshot;
     use crate::physical_optimizer::DistributedPhysicalOptimizerRule;
-    use crate::test_utils::register_parquet_tables;
+    use crate::test_utils::parquet::register_parquet_tables;
     use datafusion::error::DataFusionError;
     use datafusion::execution::SessionStateBuilder;
     use datafusion::physical_plan::displayable;

src/test_utils/insta.rs

@@ -1,10 +1,13 @@
+use datafusion::common::utils::get_available_parallelism;
 use std::env;
 
+pub use insta;
+
 #[macro_export]
 macro_rules! assert_snapshot {
     ($($arg:tt)*) => {
-        crate::test_utils::insta::settings().bind(|| {
-            insta::assert_snapshot!($($arg)*);
+        $crate::test_utils::insta::settings().bind(|| {
+            $crate::test_utils::insta::insta::assert_snapshot!($($arg)*);
         })
     };
 }
@@ -15,6 +18,10 @@ pub fn settings() -> insta::Settings {
     let cwd = env::current_dir().unwrap();
     let cwd = cwd.to_str().unwrap();
     settings.add_filter(cwd.trim_start_matches("/"), "");
+    let cpus = get_available_parallelism();
+    settings.add_filter(&format!(", {cpus}\\)"), ", CPUs)");
+    settings.add_filter(&format!("\\({cpus}\\)"), "(CPUs)");
+    settings.add_filter(&format!("input_partitions={cpus}"), "input_partitions=CPUs");
     settings.add_filter(
         r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
         "UUID",

tests/common/localhost.rs renamed to src/test_utils/localhost.rs

@@ -1,12 +1,12 @@
+use crate::{
+    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelManager, ChannelResolver, SessionBuilder,
+};
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion::{common::runtime::JoinSet, prelude::SessionConfig};
-use datafusion_distributed::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelManager, ChannelResolver, SessionBuilder,
-};
 use std::error::Error;
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;

src/test_utils/mod.rs

@@ -1,6 +1,6 @@
 pub mod insta;
-mod mock_exec;
-mod parquet;
-
-pub use mock_exec::MockExec;
-pub use parquet::register_parquet_tables;
+pub mod localhost;
+pub mod mock_exec;
+pub mod parquet;
+pub mod plan;
+pub mod tpch;

tests/common/plan.rs renamed to src/test_utils/plan.rs

@@ -1,11 +1,10 @@
+use crate::{ArrowFlightReadExec, DistributedPhysicalOptimizerRule};
 use datafusion::common::plan_err;
 use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::error::DataFusionError;
 use datafusion::physical_expr::Partitioning;
 use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::ExecutionPlan;
-use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
-use datafusion_distributed::ArrowFlightReadExec;
 use std::sync::Arc;
 
 pub fn distribute_aggregate(