apollo_integration_tests: simplify node type verifcation in restart tests (#10506)

Author: noamsp-starkware

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_restart_flow.rs

@@ -1,14 +1,12 @@
 use std::time::Duration;
 
-use apollo_deployments::deployments::consolidated::ConsolidatedNodeServiceName;
-use apollo_deployments::deployments::hybrid::HybridNodeServiceName;
+use apollo_deployments::service::NodeType;
 use apollo_infra_utils::test_utils::TestIdentifier;
 use apollo_integration_tests::integration_test_manager::{
     IntegrationTestManager,
     DEFAULT_SENDER_ACCOUNT,
 };
 use apollo_integration_tests::integration_test_utils::integration_test_setup;
-use strum::IntoEnumIterator;
 use tracing::info;
 
 #[tokio::main]
@@ -38,25 +36,9 @@ async fn main() {
     .await;
 
     // Assert that RESTART_NODE is a hybrid node.
-    assert_eq!(
-        integration_test_manager
-            .get_idle_nodes()
-            .get(&RESTART_NODE)
-            .unwrap()
-            .get_executables()
-            .len(),
-        HybridNodeServiceName::iter().count()
-    );
+    assert_eq!(integration_test_manager.get_node_type(RESTART_NODE), NodeType::Hybrid);
     // Assert that SHUTDOWN_NODE is not a consolidated node.
-    assert_ne!(
-        integration_test_manager
-            .get_idle_nodes()
-            .get(&SHUTDOWN_NODE)
-            .unwrap()
-            .get_executables()
-            .len(),
-        ConsolidatedNodeServiceName::iter().count()
-    );
+    assert_ne!(integration_test_manager.get_node_type(SHUTDOWN_NODE), NodeType::Consolidated);
 
     let mut node_indices = integration_test_manager.get_node_indices();
 

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_restart_service_single_node_flow.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use apollo_deployments::deployments::hybrid::HybridNodeServiceName;
+use apollo_deployments::service::NodeType;
 use apollo_infra_utils::test_utils::TestIdentifier;
 use apollo_integration_tests::integration_test_manager::{
     IntegrationTestManager,
@@ -38,15 +39,7 @@ async fn main() {
     .await;
 
     // Assert that RESTART_NODE is a hybrid node.
-    assert_eq!(
-        integration_test_manager
-            .get_idle_nodes()
-            .get(&RESTART_NODE)
-            .unwrap()
-            .get_executables()
-            .len(),
-        HybridNodeServiceName::iter().count()
-    );
+    assert_eq!(integration_test_manager.get_node_type(RESTART_NODE), NodeType::Hybrid);
 
     let node_indices = integration_test_manager.get_node_indices();
     integration_test_manager.run_nodes(node_indices.clone()).await;

crates/apollo_integration_tests/src/integration_test_manager.rs

@@ -372,6 +372,15 @@ impl IntegrationTestManager {
         &self.idle_nodes
     }
 
+    pub fn get_node_type(&self, node_idx: usize) -> NodeType {
+        let node_setup = self
+            .idle_nodes
+            .get(&node_idx)
+            .or_else(|| self.running_nodes.get(&node_idx).map(|node| &node.node_setup))
+            .expect("Node {node_idx} does not exist in idle or running nodes.");
+        node_setup.node_type
+    }
+
     pub fn tx_generator(&self) -> &MultiAccountTransactionGenerator {
         &self.tx_generator
     }