Skip to content

Commit b2e8623

Browse files
authored
refactor(mock/http-body): outline MockBody test body (#3611)
`MockBody` is a type that we use to implement tests for our `peek_trailers::PeekTrailersBody<B>` body middleware. this is a useful tool for mocking the polling outcomes of the inner body we wrap, which would be useful for testing other `http_body::Body` middleware. this commit moves `MockBody` out of `linkerd-http-retry`, and into a new `linkerd-mock-http-body` crate. this is added as a test dependency for the retry crate, and can now be used (rather than vendored) by other bodies. Signed-off-by: katelyn martin <[email protected]>
1 parent 40e0cf4 commit b2e8623

File tree

6 files changed

+134
-93
lines changed

6 files changed

+134
-93
lines changed

Cargo.lock

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1815,6 +1815,7 @@ dependencies = [
18151815
"linkerd-exp-backoff",
18161816
"linkerd-http-box",
18171817
"linkerd-metrics",
1818+
"linkerd-mock-http-body",
18181819
"linkerd-stack",
18191820
"linkerd-tracing",
18201821
"parking_lot",
@@ -2023,6 +2024,17 @@ dependencies = [
20232024
"tracing",
20242025
]
20252026

2027+
[[package]]
2028+
name = "linkerd-mock-http-body"
2029+
version = "0.1.0"
2030+
dependencies = [
2031+
"bytes",
2032+
"http",
2033+
"http-body",
2034+
"linkerd-error",
2035+
"tokio",
2036+
]
2037+
20262038
[[package]]
20272039
name = "linkerd-opaq-route"
20282040
version = "0.1.0"

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ members = [
4747
"linkerd/meshtls/rustls",
4848
"linkerd/meshtls/verifier",
4949
"linkerd/metrics",
50+
"linkerd/mock/http-body",
5051
"linkerd/opaq-route",
5152
"linkerd/opencensus",
5253
"linkerd/opentelemetry",

linkerd/http/retry/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,5 @@ linkerd-stack = { path = "../../stack" }
2727
[dev-dependencies]
2828
hyper = { workspace = true, features = ["deprecated"] }
2929
linkerd-tracing = { path = "../../tracing", features = ["ansi"] }
30+
linkerd-mock-http-body = { path = "../../mock/http-body" }
3031
tokio = { version = "1", features = ["macros", "rt"] }

linkerd/http/retry/src/peek_trailers.rs

Lines changed: 2 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -349,21 +349,8 @@ mod tests {
349349
use http::{HeaderMap, HeaderValue};
350350
use http_body::Body;
351351
use linkerd_error::Error;
352-
use std::{
353-
collections::VecDeque,
354-
ops::Not,
355-
pin::Pin,
356-
task::{Context, Poll},
357-
};
358-
359-
/// A "mock" body.
360-
///
361-
/// This type contains polling results for [`Body`].
362-
#[derive(Default)]
363-
struct MockBody {
364-
data_polls: VecDeque<Poll<Option<Result<Bytes, Error>>>>,
365-
trailer_polls: VecDeque<Poll<Result<Option<http::HeaderMap>, Error>>>,
366-
}
352+
use linkerd_mock_http_body::MockBody;
353+
use std::{ops::Not, task::Poll};
367354

368355
fn data() -> Option<Result<Bytes, Error>> {
369356
let bytes = Bytes::from_static(b"hello");
@@ -441,82 +428,4 @@ mod tests {
441428
assert!(peek.peek_trailers().is_none());
442429
assert!(peek.is_end_stream().not());
443430
}
444-
445-
// === impl MockBody ===
446-
447-
impl MockBody {
448-
/// Appends a poll outcome for [`Body::poll_data()`].
449-
fn then_yield_data(mut self, poll: Poll<Option<Result<Bytes, Error>>>) -> Self {
450-
self.data_polls.push_back(poll);
451-
self
452-
}
453-
454-
/// Appends a poll outcome for [`Body::poll_trailers()`].
455-
fn then_yield_trailer(
456-
mut self,
457-
poll: Poll<Result<Option<http::HeaderMap>, Error>>,
458-
) -> Self {
459-
self.trailer_polls.push_back(poll);
460-
self
461-
}
462-
463-
/// Schedules a task to be awoken.
464-
fn schedule(cx: &Context<'_>) {
465-
let waker = cx.waker().clone();
466-
tokio::spawn(async move {
467-
waker.wake();
468-
});
469-
}
470-
}
471-
472-
impl Body for MockBody {
473-
type Data = Bytes;
474-
type Error = Error;
475-
476-
fn poll_data(
477-
self: Pin<&mut Self>,
478-
cx: &mut Context<'_>,
479-
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
480-
let poll = self
481-
.get_mut()
482-
.data_polls
483-
.pop_front()
484-
.unwrap_or(Poll::Ready(None));
485-
// If we return `Poll::Pending`, we must schedule the task to be awoken.
486-
if poll.is_pending() {
487-
Self::schedule(cx);
488-
}
489-
poll
490-
}
491-
492-
fn poll_trailers(
493-
self: Pin<&mut Self>,
494-
cx: &mut Context<'_>,
495-
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
496-
let Self {
497-
data_polls,
498-
trailer_polls,
499-
} = self.get_mut();
500-
501-
let poll = if data_polls.is_empty() {
502-
trailer_polls.pop_front().unwrap_or(Poll::Ready(Ok(None)))
503-
} else {
504-
// The called has polled for trailers before exhausting the stream of DATA frames.
505-
// This indicates `PeekTrailersBody<B>` isn't respecting the contract outlined in
506-
// <https://docs.rs/http-body/0.4.6/http_body/trait.Body.html#tymethod.poll_trailers>.
507-
panic!("`poll_trailers()` was called before `poll_data()` returned `Poll::Ready(None)`");
508-
};
509-
510-
// If we return `Poll::Pending`, we must schedule the task to be awoken.
511-
if poll.is_pending() {
512-
Self::schedule(cx);
513-
}
514-
515-
poll
516-
}
517-
518-
fn is_end_stream(&self) -> bool {
519-
self.data_polls.is_empty() && self.trailer_polls.is_empty()
520-
}
521-
}
522431
}

linkerd/mock/http-body/Cargo.toml

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
[package]
2+
name = "linkerd-mock-http-body"
3+
version = "0.1.0"
4+
authors = ["Linkerd Developers <[email protected]>"]
5+
license = "Apache-2.0"
6+
edition = "2021"
7+
publish = false
8+
description = """
9+
Mock `http_body::Body` facilities for use in tests.
10+
"""
11+
12+
[dependencies]
13+
bytes = { workspace = true }
14+
http = { workspace = true }
15+
http-body = { workspace = true }
16+
linkerd-error = { path = "../../error" }
17+
tokio = { version = "1", default-features = false, features = ["rt"] }

linkerd/mock/http-body/src/lib.rs

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
//! Mock [`http_body::Body`] facilities for use in tests.
2+
//!
3+
//! See [`MockBody`] for more information.
4+
5+
use bytes::Bytes;
6+
use http_body::Body;
7+
use linkerd_error::Error;
8+
use std::{
9+
collections::VecDeque,
10+
pin::Pin,
11+
task::{Context, Poll},
12+
};
13+
14+
/// A "mock" body.
15+
///
16+
/// This type contains polling results for [`Body`].
17+
#[derive(Default)]
18+
pub struct MockBody {
19+
data_polls: VecDeque<Poll<Option<Result<Bytes, Error>>>>,
20+
trailer_polls: VecDeque<Poll<Result<Option<http::HeaderMap>, Error>>>,
21+
}
22+
23+
// === impl MockBody ===
24+
25+
impl MockBody {
26+
/// Appends a poll outcome for [`Body::poll_data()`].
27+
pub fn then_yield_data(mut self, poll: Poll<Option<Result<Bytes, Error>>>) -> Self {
28+
self.data_polls.push_back(poll);
29+
self
30+
}
31+
32+
/// Appends a poll outcome for [`Body::poll_trailers()`].
33+
pub fn then_yield_trailer(
34+
mut self,
35+
poll: Poll<Result<Option<http::HeaderMap>, Error>>,
36+
) -> Self {
37+
self.trailer_polls.push_back(poll);
38+
self
39+
}
40+
41+
/// Schedules a task to be awoken.
42+
fn schedule(cx: &Context<'_>) {
43+
let waker = cx.waker().clone();
44+
tokio::spawn(async move {
45+
waker.wake();
46+
});
47+
}
48+
}
49+
50+
impl Body for MockBody {
51+
type Data = Bytes;
52+
type Error = Error;
53+
54+
fn poll_data(
55+
self: Pin<&mut Self>,
56+
cx: &mut Context<'_>,
57+
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
58+
let poll = self
59+
.get_mut()
60+
.data_polls
61+
.pop_front()
62+
.unwrap_or(Poll::Ready(None));
63+
// If we return `Poll::Pending`, we must schedule the task to be awoken.
64+
if poll.is_pending() {
65+
Self::schedule(cx);
66+
}
67+
poll
68+
}
69+
70+
fn poll_trailers(
71+
self: Pin<&mut Self>,
72+
cx: &mut Context<'_>,
73+
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
74+
let Self {
75+
data_polls,
76+
trailer_polls,
77+
} = self.get_mut();
78+
79+
let poll = if data_polls.is_empty() {
80+
trailer_polls.pop_front().unwrap_or(Poll::Ready(Ok(None)))
81+
} else {
82+
// The caller has polled for trailers before exhausting the stream of DATA frames.
83+
// This indicates `PeekTrailersBody<B>` isn't respecting the contract outlined in
84+
// <https://docs.rs/http-body/0.4.6/http_body/trait.Body.html#tymethod.poll_trailers>.
85+
panic!(
86+
"`poll_trailers()` was called before `poll_data()` returned `Poll::Ready(None)`"
87+
);
88+
};
89+
90+
// If we return `Poll::Pending`, we must schedule the task to be awoken.
91+
if poll.is_pending() {
92+
Self::schedule(cx);
93+
}
94+
95+
poll
96+
}
97+
98+
fn is_end_stream(&self) -> bool {
99+
self.data_polls.is_empty() && self.trailer_polls.is_empty()
100+
}
101+
}

0 commit comments

Comments
 (0)