Best way to structure async integration tests

code review
1

Hey everyone! I'm writing a library that represents a node in a network (cryptographic) protocol. I want to write an integration test simulating having several nodes talking to each other and make sure it goes to completion as expected. Faking out the network interface is easy, but I'm not sure how best to structure the actual integration test. In particular I want to check that the state machine works the way it's intended and reports the right states to clients. For example:

pub enum StateKind {
    Ready,
    Voting,
    Signing,
    Complete,
}

pub trait StateChangeListener: Send + Sync {
    fn on_state_change(&self, current: StateKind, previous: StateKind);
}

pub struct Node {
    state: Mutex<StateKind>,
    state_change_listener: Option<Arc<dyn StateChangeListener>>,
}


impl Node {
    pub fn new(listener: Option<Arc<dyn StateChangeListener>>) -> Self {
        Self {
            state: Mutex::new(StateKind::Ready),
            state_change_listener: listener,
        }
    }

    pub fn state(&self) -> StateKind {
        *self.state.lock().unwrap()
    }

    fn notify_state_listener(&self, current: StateKind, previous: StateKind) {
        if let Some(listener) = self.state_change_listener.as_ref() {
            listener.on_state_change(current, previous);
        }
    }

    pub fn start_voting(&self) {
        {
            let mut guard = self.state.lock().unwrap();
            if *guard != StateKind::Ready {
                // This should actually produce an error.
                return;
            }
            *guard = StateKind::Voting;
        }
        self.notify_state_listener(StateKind::Voting, StateKind::Ready)
    }

    pub async fn process(&self) {
        // Actually process incoming messages and move through states
        // (See example code below)
    }
}

In actual practice there would be multiple Node instances talking to each other each with their own thread for processing network messages. The best I've figured out so far is to apply automock to the StateChangeListener, but this requires setting up the mock at the start of the test and doesn't allow synchronizing the mock and test function (since the mock is only checked after it goes out of scope).

::test(flavor = "multi_thread", worker_threads = 2)]
async fn node_process() {

    // Expect state changes
    let mut state_listener = MockStateChangeListener::new();
    let mut seq = mockall::Sequence::new();
    state_listener
        .expect_on_state_change()
        .with(eq(StateKind::Voting), eq(StateKind::Ready))
        .times(1)
        .in_sequence(&mut seq)
        .returning(|_, _| ());
    state_listener
        .expect_on_state_change()
        .with(eq(StateKind::Signing), eq(StateKind::Voting))
        .times(1)
        .in_sequence(&mut seq)
        .returning(|_, _| ());
    state_listener
        .expect_on_state_change()
        .with(eq(StateKind::Complete), eq(StateKind::Signing))
        .times(1)
        .in_sequence(&mut seq)
        .returning(|_, _| ());

    let node = Arc::new(Node::new(Some(Arc::new(state_listener))));
    assert_eq!(node.state(), StateKind::Ready);

    // Normally there would be more than one node communicating with each other.
    // For simplicity this is what one node looks like.
    node.start_voting();
    assert_eq!(node.state(), StateKind::Voting);
    let task_node = node.clone();
    let process_task = tokio::spawn( async move {
        task_node.process().await
    });
    // Should have a timeout on this.
    while node.state() != StateKind::Complete {
        tokio::task::yield_now().await
    }
}

Is this the cleanest way to set up an async integration test, or are there better practices I can apply? In particular, if I set up multiple nodes and have them running? And is there a way I should apply timeouts to actions / the whole test, or should I leave that to the test runner itself?

Gist with the full file.

Thanks so much!