Move all test utils to src/ and hide them behind an "integration" feature

.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,12 @@
 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)*);
         })
     };
 }

tests/common/localhost.rs → 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 → 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(

tests/custom_extension_codec.rs

@@ -1,10 +1,5 @@
-#[allow(dead_code)]
-mod common;
-
-#[cfg(test)]
+#[cfg(all(feature = "integration", test))]
 mod tests {
-    use crate::assert_snapshot;
-    use crate::common::localhost::start_localhost_context;
     use datafusion::arrow::array::Int64Array;
     use datafusion::arrow::compute::SortOptions;
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
@@ -27,9 +22,11 @@ mod tests {
     use datafusion::physical_plan::{
         displayable, execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
     };
-    use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
+    use datafusion_distributed::assert_snapshot;
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
     use datafusion_distributed::{
-        add_user_codec, with_user_codec, ArrowFlightReadExec, SessionBuilder,
+        add_user_codec, with_user_codec, ArrowFlightReadExec, DistributedPhysicalOptimizerRule,
+        SessionBuilder,
     };
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     use datafusion_proto::protobuf::proto_error;

tests/distributed_aggregation.rs

@@ -1,18 +1,18 @@
-#[allow(dead_code)]
-mod common;
-
-#[cfg(test)]
+#[cfg(all(feature = "integration", test))]
 mod tests {
-    use crate::assert_snapshot;
-    use crate::common::localhost::{start_localhost_context, NoopSessionBuilder};
-    use crate::common::parquet::register_parquet_tables;
-    use crate::common::plan::distribute_aggregate;
     use datafusion::arrow::util::pretty::pretty_format_batches;
     use datafusion::physical_plan::{displayable, execute_stream};
+    use datafusion_distributed::assert_snapshot;
+    use datafusion_distributed::test_utils::localhost::{
+        start_localhost_context, NoopSessionBuilder,
+    };
+    use datafusion_distributed::test_utils::parquet::register_parquet_tables;
+    use datafusion_distributed::test_utils::plan::distribute_aggregate;
     use futures::TryStreamExt;
     use std::error::Error;
 
     #[tokio::test]
+    #[ignore]
     async fn distributed_aggregation() -> Result<(), Box<dyn Error>> {
         // FIXME: these ports are in use on my machine, we should find unused ports
         // Changed them for now

tests/error_propagation.rs

@@ -1,9 +1,5 @@
-#[allow(dead_code)]
-mod common;
-
-#[cfg(test)]
+#[cfg(all(feature = "integration", test))]
 mod tests {
-    use crate::common::localhost::start_localhost_context;
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::error::DataFusionError;
     use datafusion::execution::{
@@ -15,9 +11,10 @@ mod tests {
     use datafusion::physical_plan::{
         execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
     };
-    use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
     use datafusion_distributed::{
-        add_user_codec, with_user_codec, ArrowFlightReadExec, SessionBuilder,
+        add_user_codec, with_user_codec, ArrowFlightReadExec, DistributedPhysicalOptimizerRule,
+        SessionBuilder,
     };
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     use datafusion_proto::protobuf::proto_error;

tests/highly_distributed_query.rs

@@ -1,14 +1,12 @@
-#[allow(dead_code)]
-mod common;
-
-#[cfg(test)]
+#[cfg(all(feature = "integration", test))]
 mod tests {
-    use crate::assert_snapshot;
-    use crate::common::localhost::{start_localhost_context, NoopSessionBuilder};
-    use crate::common::parquet::register_parquet_tables;
     use datafusion::physical_expr::Partitioning;
     use datafusion::physical_plan::{displayable, execute_stream};
-    use datafusion_distributed::ArrowFlightReadExec;
+    use datafusion_distributed::test_utils::localhost::{
+        start_localhost_context, NoopSessionBuilder,
+    };
+    use datafusion_distributed::test_utils::parquet::register_parquet_tables;
+    use datafusion_distributed::{assert_snapshot, ArrowFlightReadExec};
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;

tests/stage_planning.rs

@@ -1,16 +1,13 @@
-#[allow(dead_code)]
-mod common;
-
-#[cfg(test)]
+#[cfg(all(feature = "integration", test))]
 mod tests {
-    use crate::assert_snapshot;
-    use crate::common::tpch::tpch_query;
     use datafusion::arrow::util::pretty::pretty_format_batches;
     use datafusion::execution::SessionStateBuilder;
     use datafusion::physical_plan::{displayable, execute_stream};
     use datafusion::prelude::{SessionConfig, SessionContext};
-    use datafusion_distributed::physical_optimizer::DistributedPhysicalOptimizerRule;
-    use datafusion_distributed::stage::{display_stage_graphviz, ExecutionStage};
+    use datafusion_distributed::assert_snapshot;
+    use datafusion_distributed::test_utils::tpch::tpch_query;
+    use datafusion_distributed::DistributedPhysicalOptimizerRule;
+    use datafusion_distributed::{display_stage_graphviz, ExecutionStage};
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;