diff --git a/linkerd/http/retry/src/replay.rs b/linkerd/http/retry/src/replay.rs index e1d7c9472..5003d6c1d 100644 --- a/linkerd/http/retry/src/replay.rs +++ b/linkerd/http/retry/src/replay.rs @@ -547,8 +547,8 @@ mod tests { async fn replays_trailers() { let Test { mut tx, - mut initial, - mut replay, + initial, + replay, _trace, } = Test::new(); @@ -560,30 +560,43 @@ mod tests { tx.send_trailers(tlrs.clone()).await; drop(tx); - while initial.data().await.is_some() { - // do nothing - } - let initial_tlrs = initial.trailers().await.expect("trailers should not error"); - assert_eq!(initial_tlrs.as_ref(), Some(&tlrs)); + let read_trailers = |body: ReplayBody<_>| async move { + let mut body = crate::compat::ForwardCompatibleBody::new(body); + let _ = body + .frame() + .await + .expect("should yield a result") + .expect("should yield a frame") + .into_data() + .expect("should yield data"); + let trls = body + .frame() + .await + .expect("should yield a result") + .expect("should yield a frame") + .into_trailers() + .expect("should yield trailers"); + assert!(body.frame().await.is_none()); + trls + }; - // drop the initial body to send the data to the replay - drop(initial); + let initial_tlrs = read_trailers(initial).await; + assert_eq!(&initial_tlrs, &tlrs); - while replay.data().await.is_some() { - // do nothing - } - let replay_tlrs = replay.trailers().await.expect("trailers should not error"); - assert_eq!(replay_tlrs.as_ref(), Some(&tlrs)); + let replay_tlrs = read_trailers(replay).await; + assert_eq!(&replay_tlrs, &tlrs); } #[tokio::test] async fn trailers_only() { let Test { mut tx, - mut initial, - mut replay, + initial, + replay, _trace, } = Test::new(); + let mut initial = crate::compat::ForwardCompatibleBody::new(initial); + let mut replay = crate::compat::ForwardCompatibleBody::new(replay); let mut tlrs = HeaderMap::new(); tlrs.insert("x-hello", HeaderValue::from_str("world").unwrap()); @@ -593,16 +606,26 @@ mod tests { drop(tx); - assert!(dbg!(initial.data().await).is_none(), "no data in body"); - let initial_tlrs = initial.trailers().await.expect("trailers should not error"); - assert_eq!(initial_tlrs.as_ref(), Some(&tlrs)); + let initial_tlrs = initial + .frame() + .await + .expect("should yield a result") + .expect("should yield a frame") + .into_trailers() + .expect("should yield trailers"); + assert_eq!(&initial_tlrs, &tlrs); // drop the initial body to send the data to the replay drop(initial); - assert!(dbg!(replay.data().await).is_none(), "no data in body"); - let replay_tlrs = replay.trailers().await.expect("trailers should not error"); - assert_eq!(replay_tlrs.as_ref(), Some(&tlrs)); + let replay_tlrs = replay + .frame() + .await + .expect("should yield a result") + .expect("should yield a frame") + .into_trailers() + .expect("should yield trailers"); + assert_eq!(&replay_tlrs, &tlrs); } #[tokio::test(flavor = "current_thread")]