nanoreth/crates/net/network/tests/it/session.rs

//! Session tests

use futures::StreamExt;
use reth_eth_wire::{capability::Capability, EthVersion};
use reth_network::{
    test_utils::{PeerConfig, Testnet},
    NetworkEvent,
};
use reth_network_api::{NetworkInfo, Peers};
use reth_provider::test_utils::NoopProvider;
use std::sync::Arc;

#[tokio::test(flavor = "multi_thread")]
async fn test_session_established_with_highest_version() {
    reth_tracing::init_test_tracing();

    let net = Testnet::create(2).await;

    net.for_each(|peer| assert_eq!(0, peer.num_peers()));

    let mut handles = net.handles();
    let handle0 = handles.next().unwrap();
    let handle1 = handles.next().unwrap();
    drop(handles);

    let handle = net.spawn();

    handle0.add_peer(*handle1.peer_id(), handle1.local_addr());

    let mut events = handle0.event_listener().take(2);
    while let Some(event) = events.next().await {
        match event {
            NetworkEvent::PeerAdded(peer_id) => {
                assert_eq!(handle1.peer_id(), &peer_id);
            }
            NetworkEvent::SessionEstablished { peer_id, status, .. } => {
                assert_eq!(handle1.peer_id(), &peer_id);
                assert_eq!(status.version, EthVersion::Eth67 as u8);
            }
            _ => {
                panic!("unexpected event")
            }
        }
    }

    handle.terminate().await;
}

#[tokio::test(flavor = "multi_thread")]
async fn test_session_established_with_different_capability() {
    reth_tracing::init_test_tracing();

    let mut net = Testnet::create(1).await;

    let capabilities = vec![Capability::new("eth".into(), EthVersion::Eth66 as usize)];
    let p1 = PeerConfig::with_capabilities(Arc::new(NoopProvider::default()), capabilities);
    net.add_peer_with_config(p1).await.unwrap();

    net.for_each(|peer| assert_eq!(0, peer.num_peers()));

    let mut handles = net.handles();
    let handle0 = handles.next().unwrap();
    let handle1 = handles.next().unwrap();
    drop(handles);

    let handle = net.spawn();

    handle0.add_peer(*handle1.peer_id(), handle1.local_addr());

    let mut events = handle0.event_listener().take(2);
    while let Some(event) = events.next().await {
        match event {
            NetworkEvent::PeerAdded(peer_id) => {
                assert_eq!(handle1.peer_id(), &peer_id);
            }
            NetworkEvent::SessionEstablished { peer_id, status, .. } => {
                assert_eq!(handle1.peer_id(), &peer_id);
                assert_eq!(status.version, EthVersion::Eth66 as u8);
            }
            _ => {
                panic!("unexpected event")
            }
        }
    }

    handle.terminate().await;
}