[CX_HARDENING] Enhance DA Task Test Coverage

crates/hotshot/src/lib.rs

@@ -639,6 +639,31 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> SystemContext<T
     }
 }
 
+#[cfg(feature = "hotshot-testing")]
+impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> SystemContext<TYPES, I, V> {
+    /// Creates a lightweight version of the system handle for task state testing.
+    ///
+    /// This method provides a minimal context for task state tests, omitting the full
+    /// consensus and network task launches. It results in fewer traces and simplifies debugging.
+    ///
+    /// # Returns
+    /// A `SystemContextHandle<TYPES, I, V>` with minimal setup for task state testing.
+    pub fn create_lean_test_handle(&self) -> SystemContextHandle<TYPES, I, V> {
+        let (internal_sender, internal_receiver) = broadcast(EVENT_CHANNEL_SIZE);
+
+        SystemContextHandle {
+            consensus_registry: ConsensusTaskRegistry::new(),
+            network_registry: NetworkTaskRegistry::new(),
+            output_event_stream: self.external_event_stream.clone(),
+            internal_event_stream: (internal_sender, internal_receiver.deactivate()),
+            hotshot: self.clone().into(),
+            storage: Arc::clone(&self.storage),
+            network: Arc::clone(&self.network),
+            memberships: Arc::clone(&self.memberships),
+        }
+    }
+}
+
 /// An async broadcast channel
 type Channel<S> = (Sender<Arc<S>>, Receiver<Arc<S>>);
 

crates/testing/src/helpers.rs

@@ -5,15 +5,17 @@
 // along with the HotShot repository. If not, see <https://mit-license.org/>.
 
 #![allow(clippy::panic)]
-use std::{fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
+use std::{fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc, time::Duration};
 
 use async_broadcast::{Receiver, Sender};
+use async_compatibility_layer::art::async_timeout;
 use bitvec::bitvec;
 use committable::Committable;
 use ethereum_types::U256;
+use futures::StreamExt;
 use hotshot::{
     traits::{NodeImplementation, TestableNodeImplementation},
-    types::{BLSPubKey, SignatureKey, SystemContextHandle},
+    types::{BLSPubKey, Event, SignatureKey, SystemContextHandle},
     HotShotInitializer, Memberships, SystemContext,
 };
 use hotshot_example_types::{
@@ -40,11 +42,15 @@ use hotshot_types::{
     utils::{View, ViewInner},
     vid::{vid_scheme, VidCommitment, VidSchemeType},
     vote::{Certificate, HasViewNumber, Vote},
+    PeerConfig,
 };
 use jf_vid::VidScheme;
 use serde::Serialize;
 
-use crate::test_builder::TestDescription;
+use crate::{
+    predicates::PredicateResult, script::ExternalEventsExpectations, test_builder::TestDescription,
+    test_launcher::TestLauncher,
+};
 
 /// create the [`SystemContextHandle`] from a node id
 /// # Panics
@@ -128,6 +134,105 @@ pub async fn build_system_handle<
     .expect("Could not init hotshot")
 }
 
+/// create the [`SystemContextHandle`] from a launcher
+/// # Panics
+/// if cannot create a [`HotShotInitializer`]
+pub async fn build_system_handle_from_launcher<
+    TYPES: NodeType<InstanceState = TestInstanceState>,
+    I: NodeImplementation<
+            TYPES,
+            Storage = TestStorage<TYPES>,
+            AuctionResultsProvider = TestAuctionResultsProvider<TYPES>,
+        > + TestableNodeImplementation<TYPES>,
+    V: Versions,
+>(
+    node_id: u64,
+    launcher: TestLauncher<TYPES, I, V>,
+    membership_config: Option<Memberships<TYPES>>,
+) -> Result<SystemContextHandle<TYPES, I, V>, anyhow::Error> {
+    tracing::info!("Creating system handle for node {}", node_id);
+
+    let network = (launcher.resource_generator.channel_generator)(node_id).await;
+    let storage = (launcher.resource_generator.storage)(node_id);
+    let marketplace_config = (launcher.resource_generator.marketplace_config)(node_id);
+    let config = launcher.resource_generator.config.clone();
+
+    let known_nodes_with_stake = config.known_nodes_with_stake.clone();
+    let private_key = config.my_own_validator_config.private_key.clone();
+    let public_key = config.my_own_validator_config.public_key.clone();
+
+    let memberships = membership_config.unwrap_or_else(|| {
+        let create_membership = |nodes: &[PeerConfig<TYPES::SignatureKey>], topic: Topic| {
+            TYPES::Membership::create_election(
+                known_nodes_with_stake.clone(),
+                nodes.to_vec(),
+                topic,
+                config.fixed_leader_for_gpuvid,
+            )
+        };
+
+        Memberships {
+            quorum_membership: create_membership(&known_nodes_with_stake, Topic::Global),
+            da_membership: create_membership(&config.known_da_nodes, Topic::Da),
+            vid_membership: create_membership(&known_nodes_with_stake, Topic::Global),
+            view_sync_membership: create_membership(&known_nodes_with_stake, Topic::Global),
+        }
+    });
+
+    let initializer = HotShotInitializer::<TYPES>::from_genesis(TestInstanceState::new(
+        launcher.metadata.async_delay_config,
+    ))
+    .await
+    .unwrap();
+
+    let system_context = SystemContext::new(
+        public_key,
+        private_key,
+        node_id,
+        config,
+        memberships,
+        network,
+        initializer,
+        ConsensusMetricsValue::default(),
+        storage,
+        marketplace_config,
+    );
+
+    let system_context_handle = system_context.create_lean_test_handle();
+
+    tracing::info!("Successfully created system handle for node {}", node_id);
+    Ok(system_context_handle)
+}
+
+/// create certificate
+/// # Panics
+/// if we fail to sign the data
+pub fn build_da_cert<
+    TYPES: NodeType<SignatureKey = BLSPubKey>,
+    DATAType: Committable + Clone + Eq + Hash + Serialize + Debug + 'static,
+    VOTE: Vote<TYPES, Commitment = DATAType>,
+    CERT: Certificate<TYPES, Voteable = VOTE::Commitment>,
+>(
+    data: DATAType,
+    membership: &TYPES::Membership,
+    view: TYPES::Time,
+    public_key: &TYPES::SignatureKey,
+    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
+) -> CERT {
+    let real_qc_sig =
+        build_da_assembled_sig::<TYPES, VOTE, CERT, DATAType>(&data, membership, view);
+
+    let vote =
+        SimpleVote::<TYPES, DATAType>::create_signed_vote(data, view, public_key, private_key)
+            .expect("Failed to sign data!");
+    CERT::create_signed_certificate(
+        vote.date_commitment(),
+        vote.date().clone(),
+        real_qc_sig,
+        vote.view_number(),
+    )
+}
+
 /// create certificate
 /// # Panics
 /// if we fail to sign the data
@@ -171,6 +276,46 @@ pub fn vid_share<TYPES: NodeType>(
         .clone()
 }
 
+/// create DA signature
+/// # Panics
+/// if fails to convert node id into keypair
+pub fn build_da_assembled_sig<TYPES, VOTE, CERT, DATAType>(
+    data: &DATAType,
+    membership: &TYPES::Membership,
+    view: TYPES::Time,
+) -> <TYPES::SignatureKey as SignatureKey>::QcType
+where
+    TYPES: NodeType<SignatureKey = BLSPubKey>,
+    VOTE: Vote<TYPES>,
+    CERT: Certificate<TYPES, Voteable = VOTE::Commitment>,
+    DATAType: Committable + Clone + Eq + Hash + Serialize + Debug + 'static,
+{
+    let stake_table = membership.committee_qc_stake_table();
+    let total_nodes = stake_table.len();
+    let threshold = CERT::threshold(membership) as usize;
+    let real_qc_pp =
+        TYPES::SignatureKey::public_parameter(stake_table.clone(), U256::from(threshold as u64));
+
+    let mut signers = bitvec![0; total_nodes];
+    let sig_lists: Vec<_> = (0..threshold)
+        .map(|node_id| {
+            let (private_key, public_key) = key_pair_for_id(node_id as u64);
+            let vote = SimpleVote::<TYPES, DATAType>::create_signed_vote(
+                data.clone(),
+                view,
+                &public_key,
+                &private_key,
+            )
+            .expect("Failed to sign data");
+
+            signers.set(node_id, true);
+            vote.signature()
+        })
+        .collect();
+
+    TYPES::SignatureKey::assemble(&real_qc_pp, signers.as_bitslice(), &sig_lists)
+}
+
 /// create signature
 /// # Panics
 /// if fails to convert node id into keypair
@@ -330,7 +475,7 @@ pub fn build_da_certificate(
         payload_commit: da_payload_commitment,
     };
 
-    build_cert::<TestTypes, DaData, DaVote<TestTypes>, DaCertificate<TestTypes>>(
+    build_da_cert::<TestTypes, DaData, DaVote<TestTypes>, DaCertificate<TestTypes>>(
         da_data,
         da_membership,
         view_number,
@@ -393,3 +538,40 @@ pub fn build_fake_view_with_leaf_and_state(
         },
     }
 }
+
+pub async fn check_external_events<TYPES: NodeType, S: StreamExt<Item = Event<TYPES>> + Unpin>(
+    mut output_stream: S,
+    expectations: &[ExternalEventsExpectations<TYPES>],
+    timeout: Duration,
+) -> Result<(), String> {
+    let mut external_event_expectations_met = vec![false; expectations.len()];
+
+    while let Ok(Some(ext_event_received_output)) =
+        async_timeout(timeout, output_stream.next()).await
+    {
+        tracing::debug!("Test received Ext Event: {:?}", ext_event_received_output);
+        for (index, expectation) in expectations.iter().enumerate() {
+            if !external_event_expectations_met[index] {
+                for predicate in &expectation.output_asserts {
+                    if predicate
+                        .evaluate(&Arc::new(ext_event_received_output.clone()))
+                        .await
+                        == PredicateResult::Pass
+                    {
+                        external_event_expectations_met[index] = true;
+                        break;
+                    }
+                }
+            }
+        }
+        if external_event_expectations_met.iter().all(|&x| x) {
+            return Ok(());
+        }
+    }
+
+    if external_event_expectations_met.iter().all(|&x| x) {
+        Ok(())
+    } else {
+        Err("Not all external event expectations were met".to_string())
+    }
+}

crates/testing/src/predicates/event.rs

@@ -11,10 +11,14 @@ use async_trait::async_trait;
 use hotshot_task_impls::events::{HotShotEvent, HotShotEvent::*};
 use hotshot_types::{
     data::null_block,
+    event::Event,
     traits::{block_contents::BlockHeader, node_implementation::NodeType},
 };
 
-use crate::predicates::{Predicate, PredicateResult};
+use crate::{
+    predicates::{Predicate, PredicateResult},
+    script::ExternalEventsExpectations,
+};
 
 type EventCallback<TYPES> = Arc<dyn Fn(Arc<HotShotEvent<TYPES>>) -> bool + Send + Sync>;
 
@@ -294,3 +298,47 @@ where
     });
     Box::new(EventPredicate { check, info })
 }
+
+// New predicate type for EventType
+type ExternalEventCheckFn<TYPES> = dyn Fn(&Event<TYPES>) -> bool + Send + Sync;
+
+pub struct ExternalEventPredicate<TYPES: NodeType> {
+    check: Arc<ExternalEventCheckFn<TYPES>>,
+    info: String,
+}
+
+impl<TYPES: NodeType> std::fmt::Debug for ExternalEventPredicate<TYPES> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.info)
+    }
+}
+
+#[async_trait]
+impl<TYPES: NodeType + Send + Sync + 'static> Predicate<Arc<Event<TYPES>>>
+    for ExternalEventPredicate<TYPES>
+{
+    async fn evaluate(&self, input: &Arc<Event<TYPES>>) -> PredicateResult {
+        PredicateResult::from((self.check)(input))
+    }
+
+    async fn info(&self) -> String {
+        self.info.clone()
+    }
+}
+
+pub fn ext_event_exact<TYPES: NodeType>(
+    event: Event<TYPES>,
+) -> Box<dyn Predicate<Arc<Event<TYPES>>>> {
+    let info = format!("{:?}", event);
+    let check = Arc::new(move |e: &Event<TYPES>| {
+        e.event == event.event && e.view_number == event.view_number
+    });
+
+    Box::new(ExternalEventPredicate { check, info })
+}
+
+pub fn expect_external_events<TYPES: NodeType>(
+    predicates: Vec<Box<dyn Predicate<Arc<Event<TYPES>>>>>,
+) -> ExternalEventsExpectations<TYPES> {
+    ExternalEventsExpectations::from_outputs(predicates)
+}

crates/testing/src/script.rs

@@ -7,7 +7,7 @@
 use std::{sync::Arc, time::Duration};
 
 use hotshot_task_impls::events::HotShotEvent;
-use hotshot_types::traits::node_implementation::NodeType;
+use hotshot_types::{event::Event, traits::node_implementation::NodeType};
 
 use crate::predicates::{Predicate, PredicateResult};
 
@@ -66,6 +66,16 @@ impl<TYPES: NodeType, S> Expectations<TYPES, S> {
     }
 }
 
+pub struct ExternalEventsExpectations<TYPES: NodeType> {
+    pub output_asserts: Vec<Box<dyn Predicate<Arc<Event<TYPES>>>>>,
+}
+
+impl<TYPES: NodeType> ExternalEventsExpectations<TYPES> {
+    pub fn from_outputs(output_asserts: Vec<Box<dyn Predicate<Arc<Event<TYPES>>>>>) -> Self {
+        Self { output_asserts }
+    }
+}
+
 pub fn panic_extra_output_in_script<S>(stage_number: usize, script_name: String, output: &S)
 where
     S: std::fmt::Debug,