more tests

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -779,7 +779,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn morsel_framework_single_planner_single_morsel_no_io() -> Result<()> {
+    async fn morsel_framework_single_morsel_no_io() -> Result<()> {
         let observer = MorselObserver::new();
         let morselizer = MockMorselizer::new()
             .with_file(
@@ -816,7 +816,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn morsel_framework_single_planner_single_io_then_morsel() -> Result<()> {
+    async fn morsel_framework_single_io() -> Result<()> {
         let observer = MorselObserver::new();
         let morselizer = MockMorselizer::new()
             .with_file(
@@ -843,7 +843,9 @@ mod tests {
                 .planner_created(PlannerId(0))
                 .planner_called(PlannerId(0))
                 .io_future_created(PlannerId(0), IoFutureId(100))
+                // `with_io(IoFutureId(100), 1)` means the first poll is pending.
                 .io_future_polled(PlannerId(0), IoFutureId(100))
+                // The second poll resolves the future.
                 .io_future_polled(PlannerId(0), IoFutureId(100))
                 .io_future_resolved(PlannerId(0), IoFutureId(100))
                 .planner_called(PlannerId(0))
@@ -858,6 +860,47 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn morsel_framework_two_cpu_steps_before_morsel() -> Result<()> {
+        let observer = MorselObserver::new();
+        let morselizer = MockMorselizer::new()
+            .with_file(
+                "file1.parquet",
+                MockPlanner::builder()
+                    .with_id(PlannerId(0))
+                    .return_plan(ReturnPlanBuilder::new())
+                    .return_plan(ReturnPlanBuilder::new())
+                    .return_morsel(MorselId(10), 42)
+                    .return_none()
+                    .build(),
+            )
+            .with_observer(observer.clone());
+
+        let config = test_config(vec!["file1.parquet"]);
+        insta::assert_snapshot!(run_stream(morselizer, config).await.unwrap(), @r"
+        **** Batch: 42
+        **** Done
+        ");
+
+        assert_eq!(
+            observer.events(),
+            MorselEventsBuilder::new()
+                .morselize_file("file1.parquet")
+                .planner_created(PlannerId(0))
+                .planner_called(PlannerId(0))
+                .planner_called(PlannerId(0))
+                .planner_called(PlannerId(0))
+                .morsel_produced(PlannerId(0), MorselId(10))
+                .planner_called(PlannerId(0))
+                .morsel_stream_started(MorselId(10))
+                .morsel_stream_batch_produced(MorselId(10), BatchId(42))
+                .morsel_stream_finished(MorselId(10))
+                .build()
+        );
+
+        Ok(())
+    }
+
     fn test_config(file_names: Vec<&str>) -> FileScanConfig {
         let file_group = file_names
             .into_iter()

datafusion/datasource/src/morsel/test_utils.rs

@@ -410,6 +410,12 @@ impl ReturnPlanBuilder {
     ///
     /// `io_polls` controls how many times that future returns `Poll::Pending`
     /// before it resolves with `Poll::Ready(Ok(()))`.
+    ///
+    /// For example, `with_io(id, 1)` means:
+    /// - first poll: `Poll::Pending`
+    /// - second poll: `Poll::Ready(Ok(()))`
+    ///
+    /// So the total number of polls observed in the trace is `io_polls + 1`.
     pub fn with_io(mut self, io_future_id: IoFutureId, io_polls: usize) -> Self {
         self.io_future_id = Some(io_future_id);
         self.io_polls = io_polls;