wrappers for common usage of test_frame_writer (#2159)

This reduces boilerplate and encapsulates a good pattern of usage.
I wasn't able to make `test_frame_writer` private, but this is close
enough.

neqo-transport/src/connection/mod.rs

@@ -1146,6 +1146,19 @@ impl Connection {
         }
     }
 
+    /// A test-only output function that uses the provided writer to
+    /// pack something extra into the output.
+    #[cfg(test)]
+    pub fn test_write_frames<W>(&mut self, writer: W, now: Instant) -> Output
+    where
+        W: test_internal::FrameWriter + 'static,
+    {
+        self.test_frame_writer = Some(Box::new(writer));
+        let res = self.process_output(now);
+        self.test_frame_writer = None;
+        res
+    }
+
     /// Process input and generate output.
     #[must_use = "Output of the process function must be handled"]
     pub fn process(&mut self, dgram: Option<&Datagram>, now: Instant) -> Output {

neqo-transport/src/connection/tests/datagram.rs

@@ -399,10 +399,11 @@ fn dgram_no_allowed() {
     let mut client = default_client();
     let mut server = default_server();
     connect_force_idle(&mut client, &mut server);
-    server.test_frame_writer = Some(Box::new(InsertDatagram { data: DATA_MTU }));
-    let out = server.process_output(now()).dgram().unwrap();
-    server.test_frame_writer = None;
 
+    let out = server
+        .test_write_frames(InsertDatagram { data: DATA_MTU }, now())
+        .dgram()
+        .unwrap();
     client.process_input(&out, now());
 
     assert_error(&client, &CloseReason::Transport(Error::ProtocolViolation));
@@ -415,10 +416,10 @@ fn dgram_too_big() {
     let mut server = default_server();
     connect_force_idle(&mut client, &mut server);
 
-    server.test_frame_writer = Some(Box::new(InsertDatagram { data: DATA_MTU }));
-    let out = server.process_output(now()).dgram().unwrap();
-    server.test_frame_writer = None;
-
+    let out = server
+        .test_write_frames(InsertDatagram { data: DATA_MTU }, now())
+        .dgram()
+        .unwrap();
     client.process_input(&out, now());
 
     assert_error(&client, &CloseReason::Transport(Error::ProtocolViolation));

neqo-transport/src/connection/tests/migration.rs

@@ -25,7 +25,7 @@ use super::{
 };
 use crate::{
     cid::LOCAL_ACTIVE_CID_LIMIT,
-    connection::tests::send_something_paced,
+    connection::tests::{send_something_paced, send_with_extra},
     frame::FRAME_TYPE_NEW_CONNECTION_ID,
     packet::PacketBuilder,
     path::MAX_PATH_PROBES,
@@ -810,9 +810,7 @@ fn retire_all() {
 
     let original_cid = ConnectionId::from(get_cid(&send_something(&mut client, now())));
 
-    server.test_frame_writer = Some(Box::new(RetireAll { cid_gen }));
-    let ncid = send_something(&mut server, now());
-    server.test_frame_writer = None;
+    let ncid = send_with_extra(&mut server, RetireAll { cid_gen }, now());
 
     let new_cid_before = client.stats().frame_rx.new_connection_id;
     let retire_cid_before = client.stats().frame_tx.retire_connection_id;
@@ -861,9 +859,7 @@ fn retire_prior_to_migration_failure() {
 
     // Have the server receive the probe, but separately have it decide to
     // retire all of the available connection IDs.
-    server.test_frame_writer = Some(Box::new(RetireAll { cid_gen }));
-    let retire_all = send_something(&mut server, now());
-    server.test_frame_writer = None;
+    let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now());
 
     let resp = server.process(Some(&probe), now()).dgram().unwrap();
     assert_v4_path(&resp, true);
@@ -916,9 +912,7 @@ fn retire_prior_to_migration_success() {
 
     // Have the server receive the probe, but separately have it decide to
     // retire all of the available connection IDs.
-    server.test_frame_writer = Some(Box::new(RetireAll { cid_gen }));
-    let retire_all = send_something(&mut server, now());
-    server.test_frame_writer = None;
+    let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now());
 
     let resp = server.process(Some(&probe), now()).dgram().unwrap();
     assert_v4_path(&resp, true);
@@ -956,13 +950,11 @@ fn error_on_new_path_with_no_connection_id() {
 
     let cid_gen: Rc<RefCell<dyn ConnectionIdGenerator>> =
         Rc::new(RefCell::new(CountingConnectionIdGenerator::default()));
-    server.test_frame_writer = Some(Box::new(RetireAll { cid_gen }));
-    let retire_all = send_something(&mut server, now());
+    let retire_all = send_with_extra(&mut server, RetireAll { cid_gen }, now());
 
     client.process_input(&retire_all, now());
 
-    server.test_frame_writer = Some(Box::new(GarbageWriter {}));
-    let garbage = send_something(&mut server, now());
+    let garbage = send_with_extra(&mut server, GarbageWriter {}, now());
 
     let dgram = change_path(&garbage, DEFAULT_ADDR_V4);
     client.process_input(&dgram, now());

neqo-transport/src/connection/tests/mod.rs

@@ -17,7 +17,7 @@ use neqo_common::{event::Provider, qdebug, qtrace, Datagram, Decoder, Role};
 use neqo_crypto::{random, AllowZeroRtt, AuthenticationStatus, ResumptionToken};
 use test_fixture::{fixture_init, new_neqo_qlog, now, DEFAULT_ADDR};
 
-use super::{CloseReason, Connection, ConnectionId, Output, State};
+use super::{test_internal, CloseReason, Connection, ConnectionId, Output, State};
 use crate::{
     addr_valid::{AddressValidation, ValidateAddress},
     cc::CWND_INITIAL_PKTS,
@@ -615,6 +615,17 @@ fn send_something(sender: &mut Connection, now: Instant) -> Datagram {
     send_something_with_modifier(sender, now, Some)
 }
 
+/// Send something, but add a little something extra into the output.
+fn send_with_extra<W>(sender: &mut Connection, writer: W, now: Instant) -> Datagram
+where
+    W: test_internal::FrameWriter + 'static,
+{
+    sender.test_frame_writer = Some(Box::new(writer));
+    let res = send_something_with_modifier(sender, now, Some);
+    sender.test_frame_writer = None;
+    res
+}
+
 /// Send something on a stream from `sender` through a modifier to `receiver`.
 /// Return any ACK that might result.
 fn send_with_modifier_and_receive(

neqo-transport/src/connection/tests/recovery.rs

@@ -835,9 +835,10 @@ fn ack_for_unsent() {
     let mut server = default_server();
     connect_force_idle(&mut client, &mut server);
 
-    server.test_frame_writer = Some(Box::new(AckforUnsentWriter {}));
-    let spoofed = server.process_output(now()).dgram().unwrap();
-    server.test_frame_writer = None;
+    let spoofed = server
+        .test_write_frames(AckforUnsentWriter {}, now())
+        .dgram()
+        .unwrap();
 
     // Now deliver the packet with the spoofed ACK frame
     client.process_input(&spoofed, now());