Increase unit test coverage for diskann-disk crate

diskann-disk/src/build/chunking/checkpoint/checkpoint_context.rs

@@ -81,3 +81,75 @@ impl OwnedCheckpointContext {
         self.checkpoint_manager.mark_as_invalid()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::super::NaiveCheckpointRecordManager;
+    use super::*;
+
+    #[test]
+    fn test_checkpoint_context_new() {
+        let manager = NaiveCheckpointRecordManager;
+        let context = CheckpointContext::new(&manager, WorkStage::Start, WorkStage::End);
+
+        assert_eq!(context.current_stage(), WorkStage::Start);
+    }
+
+    #[test]
+    fn test_checkpoint_context_current_stage() {
+        let manager = NaiveCheckpointRecordManager;
+        let context =
+            CheckpointContext::new(&manager, WorkStage::QuantizeFPV, WorkStage::InMemIndexBuild);
+
+        assert_eq!(context.current_stage(), WorkStage::QuantizeFPV);
+    }
+
+    #[test]
+    fn test_checkpoint_context_get_resumption_point() {
+        let manager = NaiveCheckpointRecordManager;
+        let context = CheckpointContext::new(&manager, WorkStage::Start, WorkStage::End);
+
+        let result = context.get_resumption_point();
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), Some(0));
+    }
+
+    #[test]
+    fn test_checkpoint_context_to_owned() {
+        let manager = NaiveCheckpointRecordManager;
+        let context = CheckpointContext::new(&manager, WorkStage::Start, WorkStage::End);
+
+        let owned = context.to_owned();
+        assert_eq!(owned.current_stage(), WorkStage::Start);
+    }
+
+    #[test]
+    fn test_owned_checkpoint_context_new() {
+        let manager = Box::new(NaiveCheckpointRecordManager);
+        let context = OwnedCheckpointContext::new(
+            manager,
+            WorkStage::TrainBuildQuantizer,
+            WorkStage::PartitionData,
+        );
+
+        assert_eq!(context.current_stage(), WorkStage::TrainBuildQuantizer);
+    }
+
+    #[test]
+    fn test_owned_checkpoint_context_update() {
+        let manager = Box::new(NaiveCheckpointRecordManager);
+        let mut context = OwnedCheckpointContext::new(manager, WorkStage::Start, WorkStage::End);
+
+        let result = context.update(Progress::Completed);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_owned_checkpoint_context_mark_as_invalid() {
+        let manager = Box::new(NaiveCheckpointRecordManager);
+        let mut context = OwnedCheckpointContext::new(manager, WorkStage::Start, WorkStage::End);
+
+        let result = context.mark_as_invalid();
+        assert!(result.is_ok());
+    }
+}

diskann-disk/src/build/chunking/checkpoint/checkpoint_record.rs

@@ -218,4 +218,51 @@ mod tests {
             "Legacy code should not be able to deserialize newer enum variants"
         );
     }
+
+    #[test]
+    fn test_checkpoint_record_default() {
+        let record = CheckpointRecord::default();
+        assert!(record.is_valid());
+        assert_eq!(record.get_work_stage(), WorkStage::Start);
+    }
+
+    #[test]
+    fn test_checkpoint_record_is_valid() {
+        let record = CheckpointRecord::new();
+        assert!(record.is_valid());
+
+        let invalid_record = record.mark_as_invalid();
+        assert!(!invalid_record.is_valid());
+    }
+
+    #[test]
+    fn test_get_resumption_point_with_matching_stage() {
+        let record = CheckpointRecord::new().update_progress(42);
+        let resumption = record.get_resumption_point(WorkStage::Start);
+        assert_eq!(resumption, Some(42));
+    }
+
+    #[test]
+    fn test_get_resumption_point_with_different_stage() {
+        let record = CheckpointRecord::new();
+        let resumption = record.get_resumption_point(WorkStage::QuantizeFPV);
+        assert_eq!(resumption, None);
+    }
+
+    #[test]
+    fn test_get_resumption_point_when_invalid() {
+        let record = CheckpointRecord::new()
+            .update_progress(100)
+            .mark_as_invalid();
+        let resumption = record.get_resumption_point(WorkStage::Start);
+        assert_eq!(resumption, Some(0)); // Should return 0 when invalid
+    }
+
+    #[test]
+    fn test_advance_work_type() {
+        let record = CheckpointRecord::new();
+        let advanced = record.advance_work_type(WorkStage::QuantizeFPV).unwrap();
+        assert_eq!(advanced.get_work_stage(), WorkStage::QuantizeFPV);
+        assert!(advanced.is_valid());
+    }
 }

diskann-disk/src/build/chunking/checkpoint/checkpoint_record_manager.rs

@@ -93,3 +93,40 @@ where
         Box::new(self.clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::super::NaiveCheckpointRecordManager;
+    use super::*;
+
+    #[test]
+    fn test_checkpoint_manager_ext_execute_stage_with_resumption() {
+        let mut manager = NaiveCheckpointRecordManager;
+        let mut executed = false;
+
+        let result = manager.execute_stage(
+            WorkStage::Start,
+            WorkStage::End,
+            || {
+                executed = true;
+                Ok(42)
+            },
+            || Ok(0),
+        );
+
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), 42);
+        assert!(executed);
+    }
+
+    #[test]
+    fn test_checkpoint_manager_clone_box() {
+        let manager = NaiveCheckpointRecordManager;
+        let boxed = manager.clone_box();
+
+        // The boxed version should work the same
+        let result = boxed.get_resumption_point(WorkStage::Start);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), Some(0));
+    }
+}

diskann-disk/src/build/chunking/checkpoint/naive_checkpoint_record_manager.rs

@@ -31,3 +31,51 @@ impl CheckpointManager for NaiveCheckpointRecordManager {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::super::Progress;
+    use super::*;
+
+    #[test]
+    fn test_naive_checkpoint_record_manager_default() {
+        let manager = NaiveCheckpointRecordManager;
+        // Test get_resumption_point always returns Some(0)
+        let result = manager.get_resumption_point(WorkStage::Start);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), Some(0));
+    }
+
+    #[test]
+    fn test_naive_checkpoint_record_manager_get_resumption_point() {
+        let manager = NaiveCheckpointRecordManager;
+
+        // Test with various stages
+        for stage in [WorkStage::Start, WorkStage::End, WorkStage::QuantizeFPV] {
+            let result = manager.get_resumption_point(stage);
+            assert!(result.is_ok());
+            assert_eq!(result.unwrap(), Some(0));
+        }
+    }
+
+    #[test]
+    fn test_naive_checkpoint_record_manager_update() {
+        let mut manager = NaiveCheckpointRecordManager;
+
+        // Update should always succeed
+        let result = manager.update(Progress::Completed, WorkStage::End);
+        assert!(result.is_ok());
+
+        let result = manager.update(Progress::Processed(100), WorkStage::InMemIndexBuild);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_naive_checkpoint_record_manager_mark_as_invalid() {
+        let mut manager = NaiveCheckpointRecordManager;
+
+        // mark_as_invalid should always succeed
+        let result = manager.mark_as_invalid();
+        assert!(result.is_ok());
+    }
+}

diskann-disk/src/build/chunking/checkpoint/progress.rs

@@ -17,3 +17,28 @@ impl Progress {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_progress_map_processed() {
+        let progress = Progress::Processed(10);
+        let mapped = progress.map(|n| n * 2);
+        match mapped {
+            Progress::Processed(n) => assert_eq!(n, 20),
+            _ => panic!("Expected Processed variant"),
+        }
+    }
+
+    #[test]
+    fn test_progress_map_completed() {
+        let progress = Progress::Completed;
+        let mapped = progress.map(|n| n * 2);
+        match mapped {
+            Progress::Completed => {}
+            _ => panic!("Expected Completed variant"),
+        }
+    }
+}

diskann-disk/src/build/chunking/checkpoint/work_type.rs

@@ -23,3 +23,16 @@ pub enum WorkStage {
     Start,
     // Always add new stages at the end of the enum to avoid breaking the serialization order.
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_work_stage_serialization() {
+        let stage = WorkStage::BuildIndicesOnShards(42);
+        let serialized = bincode::serialize(&stage).unwrap();
+        let deserialized: WorkStage = bincode::deserialize(&serialized).unwrap();
+        assert_eq!(stage, deserialized);
+    }
+}

diskann-disk/src/build/chunking/continuation/chunking_config.rs

@@ -52,3 +52,59 @@ impl fmt::Display for ChunkingConfig {
         )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_chunking_config_clone() {
+        let config = ChunkingConfig {
+            continuation_checker: Box::<NaiveContinuationTracker>::default(),
+            data_compression_chunk_vector_count: 1000,
+            inmemory_build_chunk_vector_count: 5000,
+        };
+
+        let cloned = config.clone();
+
+        assert_eq!(
+            config.data_compression_chunk_vector_count,
+            cloned.data_compression_chunk_vector_count
+        );
+        assert_eq!(
+            config.inmemory_build_chunk_vector_count,
+            cloned.inmemory_build_chunk_vector_count
+        );
+    }
+
+    #[test]
+    fn test_chunking_config_display() {
+        let config = ChunkingConfig {
+            continuation_checker: Box::<NaiveContinuationTracker>::default(),
+            data_compression_chunk_vector_count: 1234,
+            inmemory_build_chunk_vector_count: 5678,
+        };
+
+        let display_string = format!("{}", config);
+
+        assert!(display_string.contains("ChunkingConfig"));
+        assert!(display_string.contains("1234"));
+        assert!(display_string.contains("5678"));
+        assert!(display_string.contains("data_compression_chunk_vector_count"));
+        assert!(display_string.contains("inmemory_build_chunk_vector_count"));
+    }
+
+    #[test]
+    fn test_chunking_config_default() {
+        let config = ChunkingConfig::default();
+
+        assert_eq!(
+            config.data_compression_chunk_vector_count,
+            PQ_COMPRESSION_DEFAULT_CHUNK_SIZE
+        );
+        assert_eq!(
+            config.inmemory_build_chunk_vector_count,
+            PQ_DEFAULT_BATCH_SIZE
+        );
+    }
+}

diskann-disk/src/build/chunking/continuation/continuation_tracker.rs

@@ -41,3 +41,30 @@ where
         Box::new(self.clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_naive_continuation_tracker_default() {
+        let tracker = NaiveContinuationTracker::default();
+        // Verify it always returns Continue
+        match tracker.get_continuation_grant() {
+            ContinuationGrant::Continue => {}
+            _ => panic!("Expected Continue"),
+        }
+    }
+
+    #[test]
+    fn test_naive_continuation_tracker_clone_box() {
+        let tracker = NaiveContinuationTracker::default();
+        let boxed = tracker.clone_box();
+
+        // The boxed version should also return Continue
+        match boxed.get_continuation_grant() {
+            ContinuationGrant::Continue => {}
+            _ => panic!("Expected Continue"),
+        }
+    }
+}