From 5dc12df98393bbb31b0f06fd9a37d62f5f91ed81 Mon Sep 17 00:00:00 2001
From: Josh Lind <josh.lind@hotmail.com>
Date: Tue, 1 Oct 2024 14:08:48 -0400
Subject: [PATCH] [Forge] Small improvements to FullNodeRebootStressTest

---
 .../src/fullnode_reboot_stress_test.rs        | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/testsuite/testcases/src/fullnode_reboot_stress_test.rs b/testsuite/testcases/src/fullnode_reboot_stress_test.rs
index 8fdd177b5af9b..46934d2616224 100644
--- a/testsuite/testcases/src/fullnode_reboot_stress_test.rs
+++ b/testsuite/testcases/src/fullnode_reboot_stress_test.rs
@@ -10,6 +10,12 @@ use rand::{seq::SliceRandom, thread_rng};
 use std::{sync::Arc, time::Duration};
 use tokio::time::Instant;
 
+// The buffer (in seconds) at the end of the test to allow for graceful shutdown
+const END_OF_TEST_BUFFER_SECS: u64 = 60;
+
+// The wait time (in seconds) between fullnode reboots
+const WAIT_TIME_BETWEEN_REBOOTS_SECS: u64 = 10;
+
 pub struct FullNodeRebootStressTest;
 
 impl Test for FullNodeRebootStressTest {
@@ -30,8 +36,19 @@ impl NetworkLoadTest for FullNodeRebootStressTest {
         _report: &mut TestReport,
         duration: Duration,
     ) -> Result<()> {
+        // Start the test timer
         let start = Instant::now();
 
+        // Ensure the total test duration is at least as long as the buffer
+        let end_of_test_buffer = Duration::from_secs(END_OF_TEST_BUFFER_SECS);
+        if duration <= end_of_test_buffer {
+            panic!(
+                "Total test duration must be at least: {:?}! Given duration: {:?}",
+                end_of_test_buffer, duration
+            );
+        }
+
+        // Collect all the fullnodes
         let all_fullnodes = {
             swarm
                 .read()
@@ -41,7 +58,9 @@ impl NetworkLoadTest for FullNodeRebootStressTest {
                 .collect::<Vec<_>>()
         };
 
-        while start.elapsed() < duration {
+        // Reboot fullnodes until the test duration is reached
+        let test_reboot_duration = duration - end_of_test_buffer;
+        while start.elapsed() < test_reboot_duration {
             {
                 let swarm = swarm.read().await;
                 let fullnode_to_reboot = {
@@ -53,7 +72,7 @@ impl NetworkLoadTest for FullNodeRebootStressTest {
                 fullnode_to_reboot.stop().await?;
                 fullnode_to_reboot.start().await?;
             }
-            tokio::time::sleep(Duration::from_secs(10)).await;
+            tokio::time::sleep(Duration::from_secs(WAIT_TIME_BETWEEN_REBOOTS_SECS)).await;
         }
 
         Ok(())