From a1c8a34711cf5020db6d1b60959079c0fd0833bf Mon Sep 17 00:00:00 2001
From: LambdaClass <121504986+lambdaclass-user@users.noreply.github.com>
Date: Fri, 13 Jan 2023 19:40:42 -0300
Subject: [PATCH] test: add test for header stage two step download (#843)

---
 crates/interfaces/src/test_utils/headers.rs |  6 ++++
 crates/stages/src/stages/headers.rs         | 31 +++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/crates/interfaces/src/test_utils/headers.rs b/crates/interfaces/src/test_utils/headers.rs
index d22d300980..f1f0b83c88 100644
--- a/crates/interfaces/src/test_utils/headers.rs
+++ b/crates/interfaces/src/test_utils/headers.rs
@@ -162,6 +162,12 @@ impl TestHeadersClient {
         lock.extend(headers);
     }
 
+    /// Clears the set.
+    pub async fn clear(&self) {
+        let mut lock = self.responses.lock().await;
+        lock.clear();
+    }
+
     /// Set response error
     pub async fn set_error(&self, err: RequestError) {
         let mut lock = self.error.lock().await;
diff --git a/crates/stages/src/stages/headers.rs b/crates/stages/src/stages/headers.rs
index 7284f5cad1..6274a34285 100644
--- a/crates/stages/src/stages/headers.rs
+++ b/crates/stages/src/stages/headers.rs
@@ -437,6 +437,37 @@ mod tests {
         );
     }
 
+    /// Execute the stage in two steps
+    #[tokio::test]
+    async fn execute_from_previous_progress() {
+        let mut runner = HeadersTestRunner::with_linear_downloader();
+        let (stage_progress, previous_stage) = (600, 1200);
+        let input = ExecInput {
+            previous_stage: Some((PREV_STAGE_ID, previous_stage)),
+            stage_progress: Some(stage_progress),
+        };
+        let headers = runner.seed_execution(input).expect("failed to seed execution");
+        let rx = runner.execute(input);
+
+        runner.client.extend(headers.iter().rev().map(|h| h.clone().unseal())).await;
+
+        // skip `after_execution` hook for linear downloader
+        let tip = headers.last().unwrap();
+        runner.consensus.update_tip(tip.hash());
+
+        let result = rx.await.unwrap();
+        assert_matches!(result, Ok(ExecOutput { done: false, stage_progress: progress }) if progress == stage_progress);
+
+        let rx = runner.execute(input);
+
+        runner.client.clear().await;
+        runner.client.extend(headers.iter().take(101).map(|h| h.clone().unseal()).rev()).await;
+
+        let result = rx.await.unwrap();
+        assert_matches!(result, Ok(ExecOutput { done: true, stage_progress }) if stage_progress == tip.number);
+        assert!(runner.validate_execution(input, result.ok()).is_ok(), "validation failed");
+    }
+
     mod test_runner {
         use crate::{
             metrics::HeaderMetrics,