From 67b48c7eed993a0d8ea33c9e48ec6c83c8d5263a Mon Sep 17 00:00:00 2001 From: Lars Eggert Date: Tue, 26 Mar 2024 16:17:22 +0200 Subject: [PATCH] test: Let packets be modified with a closure during tests Broken out of #1678 to reduce the size of that PR. --- neqo-transport/src/connection/tests/mod.rs | 83 +++++++++++++++++++--- 1 file changed, 73 insertions(+), 10 deletions(-) diff --git a/neqo-transport/src/connection/tests/mod.rs b/neqo-transport/src/connection/tests/mod.rs index b6ce08f8d1..bd7cc3ea9a 100644 --- a/neqo-transport/src/connection/tests/mod.rs +++ b/neqo-transport/src/connection/tests/mod.rs @@ -171,11 +171,12 @@ impl crate::connection::test_internal::FrameWriter for PingWriter { } /// Drive the handshake between the client and server. -fn handshake( +fn handshake_with_modifier( client: &mut Connection, server: &mut Connection, now: Instant, rtt: Duration, + mut modifier: impl FnMut(&mut Datagram), ) -> Instant { let mut a = client; let mut b = server; @@ -212,7 +213,12 @@ fn handshake( did_ping[a.role()] = true; } assert!(had_input || output.is_some()); - input = output; + if let Some(mut d) = output { + modifier(&mut d); + input = Some(d); + } else { + input = output; + } qtrace!("handshake: t += {:?}", rtt / 2); now += rtt / 2; mem::swap(&mut a, &mut b); @@ -223,6 +229,15 @@ fn handshake( now } +fn handshake( + client: &mut Connection, + server: &mut Connection, + now: Instant, + rtt: Duration, +) -> Instant { + handshake_with_modifier(client, server, now, rtt, |_| {}) +} + fn connect_fail( client: &mut Connection, server: &mut Connection, @@ -234,11 +249,12 @@ fn connect_fail( assert_error(server, &ConnectionError::Transport(server_error)); } -fn connect_with_rtt( +fn connect_with_rtt_and_modifier( client: &mut Connection, server: &mut Connection, now: Instant, rtt: Duration, + modifier: impl FnMut(&mut Datagram), ) -> Instant { fn check_rtt(stats: &Stats, rtt: Duration) { assert_eq!(stats.rtt, rtt); @@ -246,7 +262,7 @@ fn connect_with_rtt( let n = stats.frame_rx.ack + usize::from(stats.rtt_init_guess); assert_eq!(stats.rttvar, rttvar_after_n_updates(n, rtt)); } - let now = handshake(client, server, now, rtt); + let now = handshake_with_modifier(client, server, now, rtt, modifier); assert_eq!(*client.state(), State::Confirmed); assert_eq!(*server.state(), State::Confirmed); @@ -255,6 +271,15 @@ fn connect_with_rtt( now } +fn connect_with_rtt( + client: &mut Connection, + server: &mut Connection, + now: Instant, + rtt: Duration, +) -> Instant { + connect_with_rtt_and_modifier(client, server, now, rtt, |_| {}) +} + fn connect(client: &mut Connection, server: &mut Connection) { connect_with_rtt(client, server, now(), Duration::new(0, 0)); } @@ -301,8 +326,13 @@ fn assert_idle(client: &mut Connection, server: &mut Connection, rtt: Duration, } /// Connect with an RTT and then force both peers to be idle. -fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Duration) -> Instant { - let now = connect_with_rtt(client, server, now(), rtt); +fn connect_rtt_idle_with_modifier( + client: &mut Connection, + server: &mut Connection, + rtt: Duration, + modifier: impl FnMut(&mut Datagram), +) -> Instant { + let now = connect_with_rtt_and_modifier(client, server, now(), rtt, modifier); assert_idle(client, server, rtt, now); // Drain events from both as well. _ = client.events().count(); @@ -311,8 +341,20 @@ fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Durat now } +fn connect_rtt_idle(client: &mut Connection, server: &mut Connection, rtt: Duration) -> Instant { + connect_rtt_idle_with_modifier(client, server, rtt, |_| {}) +} + +fn connect_force_idle_with_modifier( + client: &mut Connection, + server: &mut Connection, + modifier: impl FnMut(&mut Datagram), +) { + connect_rtt_idle_with_modifier(client, server, Duration::new(0, 0), modifier); +} + fn connect_force_idle(client: &mut Connection, server: &mut Connection) { - connect_rtt_idle(client, server, Duration::new(0, 0)); + connect_force_idle_with_modifier(client, server, |_| {}); } fn fill_stream(c: &mut Connection, stream: StreamId) { @@ -524,12 +566,14 @@ fn assert_full_cwnd(packets: &[Datagram], cwnd: usize) { } /// Send something on a stream from `sender` to `receiver`, maybe allowing for pacing. +/// Takes a modifier function that can be used to modify the datagram before it is sent. /// Return the resulting datagram and the new time. #[must_use] -fn send_something_paced( +fn send_something_paced_with_modifier( sender: &mut Connection, mut now: Instant, allow_pacing: bool, + mut modifier: impl FnMut(&mut Datagram), ) -> (Datagram, Instant) { let stream_id = sender.stream_create(StreamType::UniDi).unwrap(); assert!(sender.stream_send(stream_id, DEFAULT_STREAM_DATA).is_ok()); @@ -544,16 +588,35 @@ fn send_something_paced( .dgram() .expect("send_something: should have something to send") } - Output::Datagram(d) => d, + Output::Datagram(mut d) => { + modifier(&mut d); + d + } Output::None => panic!("send_something: got Output::None"), }; (dgram, now) } +fn send_something_paced( + sender: &mut Connection, + now: Instant, + allow_pacing: bool, +) -> (Datagram, Instant) { + send_something_paced_with_modifier(sender, now, allow_pacing, |_| {}) +} + +fn send_something_with_modifier( + sender: &mut Connection, + now: Instant, + modifier: impl FnMut(&mut Datagram), +) -> Datagram { + send_something_paced_with_modifier(sender, now, false, modifier).0 +} + /// Send something on a stream from `sender` to `receiver`. /// Return the resulting datagram. fn send_something(sender: &mut Connection, now: Instant) -> Datagram { - send_something_paced(sender, now, false).0 + send_something_with_modifier(sender, now, |_| {}) } /// Send something on a stream from `sender` to `receiver`.