logos-blockchain-testing/examples/doc-snippets/src/custom_workload_example_expectation.rs

use async_trait::async_trait;
use testing_framework_core::scenario::{DynError, Expectation, RunContext};

pub struct ReachabilityExpectation {
    target_idx: usize,
}

impl ReachabilityExpectation {
    pub fn new(target_idx: usize) -> Self {
        Self { target_idx }
    }
}

#[async_trait]
impl Expectation for ReachabilityExpectation {
    fn name(&self) -> &str {
        "target_reachable"
    }

    async fn evaluate(&mut self, ctx: &RunContext) -> Result<(), DynError> {
        let client = ctx
            .node_clients()
            .validator_clients()
            .get(self.target_idx)
            .ok_or_else(|| {
                Box::new(std::io::Error::new(
                    std::io::ErrorKind::Other,
                    "missing target client",
                )) as DynError
            })?;

        client
            .consensus_info()
            .await
            .map(|_| ())
            .map_err(|e| e.into())
    }
}