diff --git a/rust/connlib/tunnel/src/tests/sim_relay.rs b/rust/connlib/tunnel/src/tests/sim_relay.rs index 44a84c057..8fdd0e1d1 100644 --- a/rust/connlib/tunnel/src/tests/sim_relay.rs +++ b/rust/connlib/tunnel/src/tests/sim_relay.rs @@ -2,13 +2,13 @@ use super::sim_net::{dual_ip_stack, host, Host}; use connlib_shared::messages::RelayId; use firezone_relay::{AddressFamily, AllocationPort, ClientSocket, IpStack, PeerSocket}; use proptest::prelude::*; -use rand::rngs::StdRng; +use rand::{rngs::StdRng, SeedableRng as _}; use secrecy::SecretString; use snownet::{RelaySocket, Transmit}; use std::{ borrow::Cow, collections::HashSet, - net::{SocketAddr, SocketAddrV4, SocketAddrV6}, + net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, time::{Duration, Instant, SystemTime}, }; @@ -34,7 +34,14 @@ pub(crate) fn map_explode<'a>( } impl SimRelay { - pub(crate) fn new(sut: firezone_relay::Server) -> Self { + pub(crate) fn new(seed: u64, ip4: Option, ip6: Option) -> Self { + let sut = firezone_relay::Server::new( + IpStack::from((ip4, ip6)), + rand::rngs::StdRng::seed_from_u64(seed), + 3478, + 49152..=65535, + ); + Self { sut, allocations: Default::default(), diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs index e293f8ff2..53aeb8ece 100644 --- a/rust/connlib/tunnel/src/tests/sut.rs +++ b/rust/connlib/tunnel/src/tests/sut.rs @@ -15,14 +15,12 @@ use connlib_shared::{ messages::{ClientId, GatewayId, Interface, RelayId}, DomainName, }; -use firezone_relay::IpStack; use hickory_proto::{ op::Query, rr::{RData, Record, RecordType}, }; use hickory_resolver::lookup::Lookup; use proptest_state_machine::{ReferenceStateMachine, StateMachineTest}; -use rand::SeedableRng as _; use secrecy::ExposeSecret as _; use snownet::Transmit; use std::collections::{BTreeMap, BTreeSet, VecDeque}; @@ -86,20 +84,10 @@ impl StateMachineTest for TunnelTest { let relays = ref_state .relays .iter() - .map(|(id, relay)| { - let relay = relay.map( - |seed, ip4, ip6| { - SimRelay::new(firezone_relay::Server::new( - IpStack::from((ip4, ip6)), - rand::rngs::StdRng::seed_from_u64(seed), - 3478, - 49152..=65535, - )) - }, - debug_span!("relay", rid = %id), - ); + .map(|(rid, relay)| { + let relay = relay.map(SimRelay::new, debug_span!("relay", %rid)); - (*id, relay) + (*rid, relay) }) .collect::>();