chore(iori): add test for concat merge

Author: Yesterday17

crates/iori/src/merge/auto.rs

@@ -126,7 +126,7 @@ where
         let mut namer = DuplicateOutputFileNamer::new(self.output_file.clone());
         let mut final_outputs = Vec::new();
 
-        for (i, part_index) in all_part_indexes.into_iter().enumerate() {
+        for part_index in all_part_indexes {
             let mut tracks = Vec::new();
 
             for stream_id in &streams {
@@ -170,11 +170,7 @@ where
                 tokio::fs::create_dir_all(parent).await?;
             }
 
-            let part_output_path = if i == 0 {
-                self.output_file.clone()
-            } else {
-                namer.next_path()
-            };
+            let part_output_path = namer.next_path();
 
             let output_path = if tracks.len() == 1 {
                 let track_format = tracks[0].extension().and_then(|e| e.to_str());

crates/iori/src/merge/concat.rs

@@ -97,8 +97,6 @@ async fn concat_merge(
     });
 
     let mut namer = DuplicateOutputFileNamer::new(output_path.clone());
-    let mut current_part_index: Option<u64> = None;
-    let mut output: Option<File> = None;
 
     // We don't use trim_end here because we want to handle parts individually.
     // However, we should still skip trailing failed segments in each part.
@@ -115,22 +113,18 @@ async fn concat_merge(
         let trimmed_part_segments = trim_end(part_segments, |s| !s.success);
 
         if !trimmed_part_segments.is_empty() {
-            let path = if current_part_index.is_none() {
-                output_path.clone()
-            } else {
-                namer.next_path()
-            };
+            let path = namer.next_path();
 
             let mut out = File::create(path).await?;
             for segment in trimmed_part_segments {
                 if !segment.success {
                     out = File::create(namer.next_path()).await?;
+                    continue;
                 }
 
                 let mut reader = cache.open_reader(&segment.segment).await?;
                 tokio::io::copy(&mut reader, &mut out).await?;
             }
-            current_part_index = Some(part_index);
         }
 
         part_start = part_end;
@@ -141,6 +135,10 @@ async fn concat_merge(
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+    use crate::cache::memory::MemoryCacheSource;
+    use tokio::io::AsyncWriteExt;
+
     #[test]
     fn test_trim_end() {
         let input = [1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0];
@@ -155,4 +153,135 @@ mod tests {
         let output = super::trim_end(&input, |&x| x == 0);
         assert_eq!(output, [1, 2, 3, 0, 0, 3]);
     }
+
+    async fn create_segment(
+        cache: &MemoryCacheSource,
+        sequence: u64,
+        part_index: u64,
+        data: &[u8],
+    ) -> ConcatSegment {
+        let segment = SegmentInfo {
+            sequence,
+            part_index,
+            ..Default::default()
+        };
+        let mut writer = cache.open_writer(&segment).await.unwrap().unwrap();
+        writer.write_all(data).await.unwrap();
+        writer.shutdown().await.unwrap();
+        drop(writer);
+        ConcatSegment {
+            segment,
+            success: true,
+        }
+    }
+
+    #[tokio::test]
+    async fn test_concat_merge_basic() {
+        let cache = MemoryCacheSource::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+        let output_path = temp_dir.path().join("output.ts");
+
+        let mut segments = vec![
+            create_segment(&cache, 0, 0, b"part0_seq0").await,
+            create_segment(&cache, 1, 0, b"part0_seq1").await,
+        ];
+
+        concat_merge(&mut segments, &cache, output_path.clone())
+            .await
+            .unwrap();
+
+        // Give some time for the namer Drop to run if needed,
+        // but here it's sync and should have run.
+        let content = tokio::fs::read(&output_path).await.unwrap();
+        assert_eq!(content, b"part0_seq0part0_seq1");
+    }
+
+    #[tokio::test]
+    async fn test_concat_merge_discontinuity() {
+        let cache = MemoryCacheSource::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+        let output_path = temp_dir.path().join("output.ts");
+
+        let mut segments = vec![
+            create_segment(&cache, 0, 0, b"part0_seq0").await,
+            create_segment(&cache, 1, 0, b"part0_seq1").await,
+            create_segment(&cache, 2, 1, b"part1_seq2").await,
+            create_segment(&cache, 3, 1, b"part1_seq3").await,
+        ];
+
+        concat_merge(&mut segments, &cache, output_path.clone())
+            .await
+            .unwrap();
+
+        // Check first part
+        let output_path1 = temp_dir.path().join("output.1.ts");
+        let content1 = tokio::fs::read(&output_path1).await.unwrap();
+        assert_eq!(content1, b"part0_seq0part0_seq1");
+
+        // Check second part
+        let output_path2 = temp_dir.path().join("output.2.ts");
+        let content2 = tokio::fs::read(&output_path2).await.unwrap();
+        assert_eq!(content2, b"part1_seq2part1_seq3");
+    }
+
+    #[tokio::test]
+    async fn test_concat_merge_failure() {
+        let cache = MemoryCacheSource::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+        let output_path = temp_dir.path().join("output.ts");
+
+        let mut segments = vec![
+            create_segment(&cache, 0, 0, b"part0_seq0").await,
+            ConcatSegment {
+                segment: SegmentInfo {
+                    sequence: 1,
+                    part_index: 0,
+                    ..Default::default()
+                },
+                success: false,
+            },
+            create_segment(&cache, 2, 0, b"part0_seq2").await,
+        ];
+
+        concat_merge(&mut segments, &cache, output_path.clone())
+            .await
+            .unwrap();
+
+        // First part before failure
+        let output_path1 = temp_dir.path().join("output.1.ts");
+        let content1 = tokio::fs::read(&output_path1).await.unwrap();
+        assert_eq!(content1, b"part0_seq0");
+
+        // Second part after failure
+        let output_path2 = temp_dir.path().join("output.2.ts");
+        let content2 = tokio::fs::read(&output_path2).await.unwrap();
+        assert_eq!(content2, b"part0_seq2");
+    }
+
+    #[tokio::test]
+    async fn test_concat_merge_sorting() {
+        let cache = MemoryCacheSource::new();
+        let temp_dir = tempfile::tempdir().unwrap();
+        let output_path = temp_dir.path().join("output.ts");
+
+        // Out of order segments
+        let mut segments = vec![
+            create_segment(&cache, 1, 0, b"part0_seq1").await,
+            create_segment(&cache, 0, 0, b"part0_seq0").await,
+            create_segment(&cache, 3, 1, b"part1_seq3").await,
+            create_segment(&cache, 2, 1, b"part1_seq2").await,
+        ];
+
+        concat_merge(&mut segments, &cache, output_path.clone())
+            .await
+            .unwrap();
+
+        let output_path1 = temp_dir.path().join("output.1.ts");
+        let content1 = tokio::fs::read(&output_path1).await.unwrap();
+        assert_eq!(content1, b"part0_seq0part0_seq1");
+
+        let output_path2 = temp_dir.path().join("output.2.ts");
+        let content2 = tokio::fs::read(&output_path2).await.unwrap();
+        assert_eq!(content2, b"part1_seq2part1_seq3");
+    }
 }

crates/iori/src/merge/pipe.rs

@@ -66,11 +66,7 @@ impl PipeMerger {
         let mut stream: OrderedStream<(u64, Option<SendSegment>)> = OrderedStream::new(rx);
         let future = tokio::spawn(async move {
             let mut namer = DuplicateOutputFileNamer::new(target_path.clone());
-            let mut target = Some(
-                tokio::fs::File::create(&target_path)
-                    .await
-                    .expect("Failed to create file"),
-            );
+            let mut target: Option<tokio::fs::File> = None;
             let mut current_part_index: Option<u64> = None;
 
             while let Some((_, (part_index, segment))) = stream.next().await {