[indexer-alt-framework] Improve test coverage for ingestion/client.rs

Author: tpham-mysten

crates/sui-indexer-alt-framework/src/ingestion/client.rs

@@ -45,6 +45,7 @@ pub enum FetchError {
 
 pub type FetchResult = Result<FetchData, FetchError>;
 
+#[derive(Clone)]
 pub enum FetchData {
     Raw(Bytes),
     CheckpointData(CheckpointData),
@@ -148,8 +149,7 @@ impl IngestionClient {
                     FetchData::Raw(bytes) => {
                         self.metrics.total_ingested_bytes.inc_by(bytes.len() as u64);
                         Blob::from_bytes(&bytes).map_err(|e| {
-                            self.metrics.inc_retry(
-                                checkpoint,
+                            self.metrics.inc_permanent(
                                 "deserialization",
                                 IngestionError::DeserializationError(checkpoint, e),
                             )
@@ -214,3 +214,183 @@ impl IngestionClient {
         Ok(Arc::new(data))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use prometheus::Registry;
+    use std::collections::HashMap;
+    use std::sync::Arc;
+    use std::time::Duration;
+    use tokio::sync::Mutex;
+    use tokio_util::bytes::Bytes;
+
+    use crate::ingestion::test_utils::test_checkpoint_data;
+
+    use super::*;
+
+    /// Mock implementation of IngestionClientTrait for testing
+    #[derive(Default)]
+    struct MockIngestionClient {
+        checkpoints: Arc<Mutex<HashMap<u64, FetchData>>>,
+        transient_failures: Arc<Mutex<HashMap<u64, usize>>>,
+        not_found_failures: Arc<Mutex<HashMap<u64, usize>>>,
+    }
+
+    #[async_trait::async_trait]
+    impl IngestionClientTrait for MockIngestionClient {
+        async fn fetch(&self, checkpoint: u64) -> FetchResult {
+            // Check for not found failures
+            if let Some(remaining) = self.not_found_failures.lock().await.get_mut(&checkpoint) {
+                if *remaining > 0 {
+                    *remaining -= 1;
+                    return Err(FetchError::NotFound);
+                }
+            }
+
+            // Check for transient failures
+            if let Some(remaining) = self.transient_failures.lock().await.get_mut(&checkpoint) {
+                if *remaining > 0 {
+                    *remaining -= 1;
+                    return Err(FetchError::Transient {
+                        reason: "mock_transient_error",
+                        error: anyhow::anyhow!("Mock transient error"),
+                    });
+                }
+            }
+
+            // Return the checkpoint data if it exists
+            if let Some(data) = self.checkpoints.lock().await.get(&checkpoint) {
+                Ok(data.clone())
+            } else {
+                Err(FetchError::NotFound)
+            }
+        }
+    }
+
+    fn setup_test() -> (IngestionClient, Arc<MockIngestionClient>) {
+        let registry = Registry::new_custom(Some("test".to_string()), None).unwrap();
+        let metrics = IndexerMetrics::new(&registry);
+        let mock_client = Arc::new(MockIngestionClient::default());
+        let client = IngestionClient::new_impl(mock_client.clone(), metrics);
+        (client, mock_client)
+    }
+
+    #[tokio::test]
+    async fn test_fetch_raw_bytes_success() {
+        let (client, mock) = setup_test();
+
+        // Create test data using test_checkpoint_data
+        let bytes = Bytes::from(test_checkpoint_data(1));
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::Raw(bytes.clone()));
+
+        // Fetch and verify
+        let result = client.fetch(1).await.unwrap();
+        assert_eq!(result.checkpoint_summary.sequence_number(), &1);
+    }
+
+    #[tokio::test]
+    async fn test_fetch_checkpoint_data_success() {
+        let (client, mock) = setup_test();
+
+        // Create test data using test_checkpoint_data
+        let bytes = test_checkpoint_data(1);
+        let checkpoint_data: CheckpointData = Blob::from_bytes(&bytes).unwrap();
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::CheckpointData(checkpoint_data.clone()));
+
+        // Fetch and verify
+        let result = client.fetch(1).await.unwrap();
+        assert_eq!(result.checkpoint_summary.sequence_number(), &1);
+    }
+
+    #[tokio::test]
+    async fn test_fetch_not_found() {
+        let (client, _) = setup_test();
+
+        // Try to fetch non-existent checkpoint
+        let result = client.fetch(1).await;
+        assert!(matches!(result, Err(IngestionError::NotFound(1))));
+    }
+
+    #[tokio::test]
+    async fn test_fetch_transient_error_with_retry() {
+        let (client, mock) = setup_test();
+
+        // Create test data using test_checkpoint_data
+        let bytes = test_checkpoint_data(1);
+        let checkpoint_data: CheckpointData = Blob::from_bytes(&bytes).unwrap();
+
+        // Add checkpoint to mock with 2 transient failures
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::CheckpointData(checkpoint_data.clone()));
+        mock.transient_failures.lock().await.insert(1, 2);
+
+        // Fetch and verify it succeeds after retries
+        let result = client.fetch(1).await.unwrap();
+        assert_eq!(*result.checkpoint_summary.sequence_number(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_wait_for_checkpoint_with_retry() {
+        let (client, mock) = setup_test();
+
+        // Create test data using test_checkpoint_data
+        let bytes = test_checkpoint_data(1);
+        let checkpoint_data: CheckpointData = Blob::from_bytes(&bytes).unwrap();
+
+        // Add checkpoint to mock with 1 not_found failures
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::CheckpointData(checkpoint_data));
+        mock.not_found_failures.lock().await.insert(1, 1);
+
+        // Wait for checkpoint with short retry interval
+        let result = client.wait_for(1, Duration::from_millis(50)).await.unwrap();
+        assert_eq!(result.checkpoint_summary.sequence_number(), &1);
+    }
+
+    #[tokio::test]
+    async fn test_wait_for_checkpoint_instant() {
+        let (client, mock) = setup_test();
+
+        // Create test data using test_checkpoint_data
+        let bytes = test_checkpoint_data(1);
+        let checkpoint_data: CheckpointData = Blob::from_bytes(&bytes).unwrap();
+
+        // Add checkpoint to mock with no failures - data should be available immediately
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::CheckpointData(checkpoint_data));
+
+        // Wait for checkpoint with short retry interval
+        let result = client.wait_for(1, Duration::from_millis(50)).await.unwrap();
+        assert_eq!(result.checkpoint_summary.sequence_number(), &1);
+    }
+
+    #[tokio::test]
+    async fn test_wait_for_checkpoint_permanent_deserialization_error() {
+        let (client, mock) = setup_test();
+
+        // Add invalid data that will cause a deserialization error
+        mock.checkpoints
+            .lock()
+            .await
+            .insert(1, FetchData::Raw(Bytes::from("invalid data")));
+
+        // Wait for checkpoint should fail immediately with a permanent error
+        let result = client.wait_for(1, Duration::from_millis(50)).await;
+        assert!(matches!(
+            result,
+            Err(IngestionError::DeserializationError(1, _))
+        ));
+    }
+}

crates/sui-indexer-alt-framework/src/ingestion/remote_client.rs

@@ -207,26 +207,13 @@ pub(crate) mod tests {
         assert_eq!(42, checkpoint.checkpoint_summary.sequence_number)
     }
 
-    /// Treat deserialization failure as another kind of transient error -- all checkpoint data
-    /// that is fetched should be valid (deserializable as a `CheckpointData`).
     #[tokio::test]
-    async fn retry_on_deserialization_error() {
+    async fn fail_on_deserialization_error() {
         let server = MockServer::start().await;
-        let times: Mutex<u64> = Mutex::new(0);
-        respond_with(&server, move |_: &Request| {
-            let mut times = times.lock().unwrap();
-            *times += 1;
-            if *times < 3 {
-                status(StatusCode::OK).set_body_bytes(vec![])
-            } else {
-                status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42))
-            }
-        })
-        .await;
+        respond_with(&server, status(StatusCode::OK).set_body_bytes(vec![])).await;
 
         let client = remote_test_client(server.uri());
-        let checkpoint = client.fetch(42).await.unwrap();
-
-        assert_eq!(42, checkpoint.checkpoint_summary.sequence_number)
+        let result = client.fetch(42).await;
+        assert!(matches!(result, Err(Error::DeserializationError(42, _))));
     }
 }

crates/sui-indexer-alt-framework/src/metrics.rs

@@ -53,6 +53,7 @@ pub struct IndexerMetrics {
     pub total_ingested_bytes: IntCounter,
     pub total_ingested_transient_retries: IntCounterVec,
     pub total_ingested_not_found_retries: IntCounter,
+    pub total_ingested_permanent_errors: IntCounterVec,
 
     // Checkpoint lag metrics for the ingestion pipeline.
     pub latest_ingested_checkpoint: IntGauge,
@@ -199,6 +200,12 @@ impl IndexerMetrics {
                 registry,
             )
             .unwrap(),
+            total_ingested_permanent_errors: register_int_counter_vec_with_registry!(
+                "total_ingested_permanent_errors",
+                "Total number of permanent errors while fetching data from the remote store",
+                &["reason"],
+                registry,
+            ).unwrap(),
             latest_ingested_checkpoint: register_int_gauge_with_registry!(
                 "indexer_latest_ingested_checkpoint",
                 "Latest checkpoint sequence number fetched from the remote store",
@@ -604,6 +611,16 @@ impl IndexerMetrics {
 
         backoff::Error::transient(error)
     }
+
+    /// Register that we're NOT retrying a checkpoint fetch due to a permanent error, logging the
+    /// reason and error.
+    pub fn inc_permanent(&self, reason: &str, error: Error) -> backoff::Error<Error> {
+        self.total_ingested_permanent_errors
+            .with_label_values(&[reason])
+            .inc();
+
+        backoff::Error::permanent(error)
+    }
 }
 
 impl CheckpointLagMetricReporter {