zebra_test/mock_service.rs
1//! Some helpers to make it simpler to mock Tower services.
2//!
3//! A [`MockService`] is a generic [`tower::Service`] implementation that allows intercepting
4//! requests, responding to them individually, and checking that there are no requests to be
5//! received (at least during a period of time). The [`MockService`] can be built for proptests or
6//! for normal Rust unit tests.
7//!
8//! # Example
9//!
10//! ```
11//! use zebra_test::mock_service::MockService;
12//! # use tower::ServiceExt;
13//!
14//! # let reactor = tokio::runtime::Builder::new_current_thread()
15//! # .enable_all()
16//! # .build()
17//! # .expect("Failed to build Tokio runtime");
18//! #
19//! # reactor.block_on(async {
20//! let mut mock_service = MockService::build().for_unit_tests();
21//! let mut service = mock_service.clone();
22//! #
23//! # // Add types to satisfy the compiler's type inference for the `Error` type.
24//! # let _typed_mock_service: MockService<_, _, _> = mock_service.clone();
25//!
26//! let call = tokio::spawn(mock_service.clone().oneshot("hello"));
27//!
28//! mock_service
29//! .expect_request("hello").await
30//! .respond("hi!");
31//!
32//! mock_service.expect_no_requests().await;
33//!
34//! let response = call
35//! .await
36//! .expect("Failed to run call on the background")
37//! .expect("Failed to receive response from service");
38//!
39//! assert_eq!(response, "hi!");
40//! # });
41//! ```
42
43use std::{
44 fmt::Debug,
45 marker::PhantomData,
46 sync::{
47 atomic::{AtomicUsize, Ordering},
48 Arc,
49 },
50 task::{Context, Poll},
51 time::Duration,
52};
53
54use futures::{future::BoxFuture, FutureExt};
55use proptest::prelude::*;
56use tokio::{
57 sync::{
58 broadcast::{self, error::RecvError},
59 oneshot, Mutex,
60 },
61 time::timeout,
62};
63use tower::{BoxError, Service};
64
65/// The default size of the channel that forwards received requests.
66///
67/// If requests are received faster than the test code can consume them, some requests may be
68/// ignored.
69///
70/// This value can be configured in the [`MockService`] using
71/// [`MockServiceBuilder::with_proxy_channel_size`].
72const DEFAULT_PROXY_CHANNEL_SIZE: usize = 100;
73
74/// The default timeout before considering a request has not been received.
75///
76/// This is the time that the mocked service waits before considering a request will not be
77/// received. It can be configured in the [`MockService`] using
78/// [`MockServiceBuilder::with_max_request_delay`].
79///
80/// Note that if a test checks that no requests are received, each check has to wait for this
81/// amount of time, so this may affect the test execution time.
82///
83/// We've seen delays up to 67ms on busy Linux and macOS machines,
84/// and some other timeout failures even with a 150ms timeout.
85pub const DEFAULT_MAX_REQUEST_DELAY: Duration = Duration::from_millis(300);
86
87/// An internal type representing the item that's sent in the [`broadcast`] channel.
88///
89/// The actual type that matters is the [`ResponseSender`] but since there could be more than one
90/// [`MockService`] verifying requests, the type must be wrapped so that it can be shared by all
91/// receivers:
92///
93/// - The [`Arc`] makes sure the instance is on the heap, and can be shared properly between
94/// threads and dropped when no longer needed.
95/// - The [`Mutex`] ensures only one [`MockService`] instance can reply to the received request.
96/// - The [`Option`] forces the [`MockService`] that handles the request to take ownership of it
97/// because sending a response also forces the [`ResponseSender`] to be dropped.
98type ProxyItem<Request, Response, Error> =
99 Arc<Mutex<Option<ResponseSender<Request, Response, Error>>>>;
100
101/// A service implementation that allows intercepting requests for checking them.
102///
103/// The type is generic over the request and response types, and also has an extra generic type
104/// parameter that's used as a tag to determine if the internal assertions should panic or return
105/// errors for proptest minimization. See [`AssertionType`] for more information.
106///
107/// The mock service can be cloned, and provides methods for checking the received requests as well
108/// as responding to them individually.
109///
110/// Internally, the instance that's operating as the service will forward requests to a
111/// [`broadcast`] channel that the other instances listen to.
112///
113/// See the [module-level documentation][`super::mock_service`] for an example.
114pub struct MockService<Request, Response, Assertion, Error = BoxError> {
115 receiver: broadcast::Receiver<ProxyItem<Request, Response, Error>>,
116 sender: broadcast::Sender<ProxyItem<Request, Response, Error>>,
117 poll_count: Arc<AtomicUsize>,
118 max_request_delay: Duration,
119 _assertion_type: PhantomData<Assertion>,
120}
121
122/// A builder type to create a [`MockService`].
123///
124/// Allows changing specific parameters used by the [`MockService`], if necessary. The default
125/// parameters should be reasonable for most cases.
126#[derive(Default)]
127pub struct MockServiceBuilder {
128 proxy_channel_size: Option<usize>,
129 max_request_delay: Option<Duration>,
130}
131
132/// A helper type for responding to incoming requests.
133///
134/// An instance of this type is created for each request received by the [`MockService`]. It
135/// contains the received request and a [`oneshot::Sender`] that can be used to respond to the
136/// request.
137///
138/// If a response is not sent, the channel is closed and a [`BoxError`] is returned by the service
139/// to the caller that sent the request.
140#[must_use = "Tests may fail if a response is not sent back to the caller"]
141pub struct ResponseSender<Request, Response, Error> {
142 request: Request,
143 response_sender: oneshot::Sender<Result<Response, Error>>,
144}
145
146/// The [`tower::Service`] implementation of the [`MockService`].
147///
148/// The [`MockService`] is always ready, and it intercepts the requests wrapping them in a
149/// [`ResponseSender`] which can be used to send a response.
150impl<Request, Response, Assertion, Error> Service<Request>
151 for MockService<Request, Response, Assertion, Error>
152where
153 Request: Send + 'static,
154 Response: Send + 'static,
155 Error: Send + 'static,
156{
157 type Response = Response;
158 type Error = Error;
159 type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
160
161 fn poll_ready(&mut self, _context: &mut Context) -> Poll<Result<(), Self::Error>> {
162 self.poll_count.fetch_add(1, Ordering::SeqCst);
163 Poll::Ready(Ok(()))
164 }
165
166 fn call(&mut self, request: Request) -> Self::Future {
167 let (response_sender, response_receiver) = ResponseSender::new(request);
168 let proxy_item = Arc::new(Mutex::new(Some(response_sender)));
169
170 let _ = self.sender.send(proxy_item);
171
172 response_receiver
173 .map(|response| {
174 response.expect("A response was not sent by the `MockService` for a request")
175 })
176 .boxed()
177 }
178}
179
180/// An entry point for starting the [`MockServiceBuilder`].
181///
182/// This `impl` block exists for ergonomic reasons. The generic type parameters don't matter,
183/// because they are actually set by [`MockServiceBuilder::finish`].
184impl MockService<(), (), ()> {
185 /// Create a [`MockServiceBuilder`] to help with the creation of a [`MockService`].
186 pub fn build() -> MockServiceBuilder {
187 MockServiceBuilder::default()
188 }
189}
190
191impl MockServiceBuilder {
192 /// Configure the size of the proxy channel used for sending intercepted requests.
193 ///
194 /// This determines the maximum amount of requests that are kept in queue before the oldest
195 /// request is dropped. This means that any tests that receive too many requests might ignore
196 /// some requests if this parameter isn't properly configured.
197 ///
198 /// The default value of 100 should be enough for most cases.
199 ///
200 /// # Example
201 ///
202 /// ```
203 /// # use zebra_test::mock_service::MockService;
204 /// #
205 /// let mock_service = MockService::build()
206 /// .with_proxy_channel_size(100)
207 /// .for_prop_tests();
208 /// #
209 /// # // Add types to satisfy the compiler's type inference.
210 /// # let typed_mock_service: MockService<(), (), _> = mock_service;
211 /// ```
212 pub fn with_proxy_channel_size(mut self, size: usize) -> Self {
213 self.proxy_channel_size = Some(size);
214 self
215 }
216
217 /// Configure the time to wait for a request before considering no requests will be received.
218 ///
219 /// This determines the maximum amount of time that the [`MockService`] will wait for a request
220 /// to be received before considering that a request will not be received.
221 ///
222 /// The default value of 25 ms should be enough for most cases.
223 ///
224 /// # Example
225 ///
226 /// ```
227 /// # use std::time::Duration;
228 /// #
229 /// # use zebra_test::mock_service::MockService;
230 /// #
231 /// let mock_service = MockService::build()
232 /// .with_max_request_delay(Duration::from_millis(25))
233 /// .for_unit_tests();
234 /// #
235 /// # // Add types to satisfy the compiler's type inference.
236 /// # let typed_mock_service: MockService<(), (), _> = mock_service;
237 /// ```
238 pub fn with_max_request_delay(mut self, max_request_delay: Duration) -> Self {
239 self.max_request_delay = Some(max_request_delay);
240 self
241 }
242
243 /// Create a [`MockService`] to be used in [`mod@proptest`]s.
244 ///
245 /// The assertions performed by [`MockService`] use the macros provided by [`mod@proptest`], like
246 /// [`prop_assert`].
247 pub fn for_prop_tests<Request, Response, Error>(
248 self,
249 ) -> MockService<Request, Response, PropTestAssertion, Error> {
250 self.finish()
251 }
252
253 /// Create a [`MockService`] to be used in Rust unit tests.
254 ///
255 /// The assertions performed by [`MockService`] use the macros provided by default in Rust,
256 /// like [`assert`].
257 pub fn for_unit_tests<Request, Response, Error>(
258 self,
259 ) -> MockService<Request, Response, PanicAssertion, Error> {
260 self.finish()
261 }
262
263 /// An internal helper method to create the actual [`MockService`].
264 ///
265 /// Note that this is used by both [`Self::for_prop_tests`] and [`Self::for_unit_tests`], the
266 /// only difference being the `Assertion` generic type parameter, which Rust infers
267 /// automatically.
268 pub fn finish<Request, Response, Assertion, Error>(
269 self,
270 ) -> MockService<Request, Response, Assertion, Error> {
271 let proxy_channel_size = self
272 .proxy_channel_size
273 .unwrap_or(DEFAULT_PROXY_CHANNEL_SIZE);
274 let (sender, receiver) = broadcast::channel(proxy_channel_size);
275
276 MockService {
277 receiver,
278 sender,
279 poll_count: Arc::new(AtomicUsize::new(0)),
280 max_request_delay: self.max_request_delay.unwrap_or(DEFAULT_MAX_REQUEST_DELAY),
281 _assertion_type: PhantomData,
282 }
283 }
284}
285
286/// Implementation of [`MockService`] methods that use standard Rust panicking assertions.
287impl<Request, Response, Error> MockService<Request, Response, PanicAssertion, Error> {
288 /// Expect a specific request to be received.
289 ///
290 /// The expected request should be the next one in the internal queue, or if the queue is
291 /// empty, it should be received in at most the max delay time configured by
292 /// [`MockServiceBuilder::with_max_request_delay`].
293 ///
294 /// If the received request matches the expected request, a [`ResponseSender`] is returned
295 /// which can be used to inspect the request and respond to it. If no response is sent, the
296 /// sender of the requests receives an error.
297 ///
298 /// # Panics
299 ///
300 /// If no request is received or if a request is received that's not equal to the expected
301 /// request, this method panics.
302 ///
303 /// # Example
304 ///
305 /// ```
306 /// # use zebra_test::mock_service::MockService;
307 /// # use tower::ServiceExt;
308 /// #
309 /// # let reactor = tokio::runtime::Builder::new_current_thread()
310 /// # .enable_all()
311 /// # .build()
312 /// # .expect("Failed to build Tokio runtime");
313 /// #
314 /// # reactor.block_on(async {
315 /// # let mut mock_service: MockService<_, _, _> = MockService::build().for_unit_tests();
316 /// # let mut service = mock_service.clone();
317 /// #
318 /// let call = tokio::spawn(mock_service.clone().oneshot("request"));
319 ///
320 /// mock_service.expect_request("request").await.respond("response");
321 ///
322 /// assert!(matches!(call.await, Ok(Ok("response"))));
323 /// # });
324 /// ```
325 pub async fn expect_request(
326 &mut self,
327 expected: Request,
328 ) -> ResponseSender<Request, Response, Error>
329 where
330 Request: PartialEq + Debug,
331 {
332 let response_sender = self.next_request().await;
333
334 assert_eq!(
335 response_sender.request,
336 expected,
337 "received an unexpected request\n \
338 in {}",
339 std::any::type_name::<Self>(),
340 );
341
342 response_sender
343 }
344
345 /// Expect a request to be received that matches a specified condition.
346 ///
347 /// There should be a request already in the internal queue, or a request should be received in
348 /// at most the max delay time configured by [`MockServiceBuilder::with_max_request_delay`].
349 ///
350 /// The received request is passed to the `condition` function, which should return `true` if
351 /// it matches the expected condition or `false` otherwise. If `true` is returned, a
352 /// [`ResponseSender`] is returned which can be used to inspect the request again and respond
353 /// to it. If no response is sent, the sender of the requests receives an error.
354 ///
355 /// # Panics
356 ///
357 /// If the `condition` function returns `false`, this method panics.
358 ///
359 /// # Example
360 ///
361 /// ```
362 /// # use zebra_test::mock_service::MockService;
363 /// # use tower::ServiceExt;
364 /// #
365 /// # let reactor = tokio::runtime::Builder::new_current_thread()
366 /// # .enable_all()
367 /// # .build()
368 /// # .expect("Failed to build Tokio runtime");
369 /// #
370 /// # reactor.block_on(async {
371 /// # let mut mock_service: MockService<_, _, _> = MockService::build().for_unit_tests();
372 /// # let mut service = mock_service.clone();
373 /// #
374 /// let call = tokio::spawn(mock_service.clone().oneshot(1));
375 ///
376 /// mock_service.expect_request_that(|request| *request > 0).await.respond("response");
377 ///
378 /// assert!(matches!(call.await, Ok(Ok("response"))));
379 /// # });
380 /// ```
381 pub async fn expect_request_that(
382 &mut self,
383 condition: impl FnOnce(&Request) -> bool,
384 ) -> ResponseSender<Request, Response, Error>
385 where
386 Request: Debug,
387 {
388 let response_sender = self.next_request().await;
389
390 assert!(
391 condition(&response_sender.request),
392 "condition was false for request: {:?},\n \
393 in {}",
394 response_sender.request,
395 std::any::type_name::<Self>(),
396 );
397
398 response_sender
399 }
400
401 /// Expect no requests to be received.
402 ///
403 /// The internal queue of received requests should be empty, and no new requests should arrive
404 /// for the max delay time configured by [`MockServiceBuilder::with_max_request_delay`].
405 ///
406 /// # Panics
407 ///
408 /// If the queue is not empty or if a request is received before the max request delay timeout
409 /// expires.
410 ///
411 /// # Example
412 ///
413 /// ```
414 /// # use zebra_test::mock_service::MockService;
415 /// # use tower::ServiceExt;
416 /// #
417 /// # let reactor = tokio::runtime::Builder::new_current_thread()
418 /// # .enable_all()
419 /// # .build()
420 /// # .expect("Failed to build Tokio runtime");
421 /// #
422 /// # reactor.block_on(async {
423 /// # let mut mock_service: MockService<(), (), _> = MockService::build().for_unit_tests();
424 /// #
425 /// mock_service.expect_no_requests().await;
426 /// # });
427 /// ```
428 pub async fn expect_no_requests(&mut self)
429 where
430 Request: Debug,
431 {
432 if let Some(response_sender) = self.try_next_request().await {
433 panic!(
434 "received an unexpected request: {:?},\n \
435 in {}",
436 response_sender.request,
437 std::any::type_name::<Self>(),
438 );
439 }
440 }
441
442 /// Returns the next request from the queue,
443 /// or panics if there are no requests after a short timeout.
444 ///
445 /// Returns the next request in the internal queue or waits at most the max delay time
446 /// configured by [`MockServiceBuilder::with_max_request_delay`] for a new request to be
447 /// received, and then returns that.
448 ///
449 /// # Panics
450 ///
451 /// If the queue is empty and a request is not received before the max request delay timeout
452 /// expires.
453 async fn next_request(&mut self) -> ResponseSender<Request, Response, Error> {
454 match self.try_next_request().await {
455 Some(request) => request,
456 None => panic!(
457 "timeout while waiting for a request\n \
458 in {}",
459 std::any::type_name::<Self>(),
460 ),
461 }
462 }
463
464 /// Returns a count of the number of times this service has been polled.
465 ///
466 /// Note: The poll count wraps around on overflow.
467 pub fn poll_count(&self) -> usize {
468 self.poll_count.load(Ordering::SeqCst)
469 }
470}
471
472/// Implementation of [`MockService`] methods that use [`mod@proptest`] assertions.
473impl<Request, Response, Error> MockService<Request, Response, PropTestAssertion, Error> {
474 /// Expect a specific request to be received.
475 ///
476 /// The expected request should be the next one in the internal queue, or if the queue is
477 /// empty, it should be received in at most the max delay time configured by
478 /// [`MockServiceBuilder::with_max_request_delay`].
479 ///
480 /// If the received request matches the expected request, a [`ResponseSender`] is returned
481 /// which can be used to inspect the request and respond to it. If no response is sent, the
482 /// sender of the requests receives an error.
483 ///
484 /// If no request is received or if a request is received that's not equal to the expected
485 /// request, this method returns an error generated by a [`mod@proptest`] assertion.
486 ///
487 /// # Example
488 ///
489 /// ```
490 /// # use proptest::prelude::*;
491 /// # use tower::ServiceExt;
492 /// #
493 /// # use zebra_test::mock_service::MockService;
494 /// #
495 /// # let reactor = tokio::runtime::Builder::new_current_thread()
496 /// # .enable_all()
497 /// # .build()
498 /// # .expect("Failed to build Tokio runtime");
499 /// #
500 /// # reactor.block_on(async {
501 /// # let test_code = || async {
502 /// # let mut mock_service: MockService<_, _, _> =
503 /// # MockService::build().for_prop_tests();
504 /// # let mut service = mock_service.clone();
505 /// #
506 /// let call = tokio::spawn(mock_service.clone().oneshot("request"));
507 ///
508 /// // NOTE: The try operator `?` is required for errors to be handled by proptest.
509 /// mock_service
510 /// .expect_request("request").await?
511 /// .respond("response");
512 ///
513 /// prop_assert!(matches!(call.await, Ok(Ok("response"))));
514 /// #
515 /// # Ok::<(), TestCaseError>(())
516 /// # };
517 /// # test_code().await
518 /// # }).unwrap();
519 /// ```
520 pub async fn expect_request(
521 &mut self,
522 expected: Request,
523 ) -> Result<ResponseSender<Request, Response, Error>, TestCaseError>
524 where
525 Request: PartialEq + Debug,
526 {
527 let response_sender = self.next_request().await?;
528
529 prop_assert_eq!(
530 &response_sender.request,
531 &expected,
532 "received an unexpected request\n \
533 in {}",
534 std::any::type_name::<Self>(),
535 );
536
537 Ok(response_sender)
538 }
539
540 /// Expect a request to be received that matches a specified condition.
541 ///
542 /// There should be a request already in the internal queue, or a request should be received in
543 /// at most the max delay time configured by [`MockServiceBuilder::with_max_request_delay`].
544 ///
545 /// The received request is passed to the `condition` function, which should return `true` if
546 /// it matches the expected condition or `false` otherwise. If `true` is returned, a
547 /// [`ResponseSender`] is returned which can be used to inspect the request again and respond
548 /// to it. If no response is sent, the sender of the requests receives an error.
549 ///
550 /// If the `condition` function returns `false`, this method returns an error generated by a
551 /// [`mod@proptest`] assertion.
552 ///
553 /// # Example
554 ///
555 /// ```
556 /// # use proptest::prelude::*;
557 /// # use tower::ServiceExt;
558 /// #
559 /// # use zebra_test::mock_service::MockService;
560 /// #
561 /// # let reactor = tokio::runtime::Builder::new_current_thread()
562 /// # .enable_all()
563 /// # .build()
564 /// # .expect("Failed to build Tokio runtime");
565 /// #
566 /// # reactor.block_on(async {
567 /// # let test_code = || async {
568 /// # let mut mock_service: MockService<_, _, _> =
569 /// # MockService::build().for_prop_tests();
570 /// # let mut service = mock_service.clone();
571 /// #
572 /// let call = tokio::spawn(mock_service.clone().oneshot(1));
573 ///
574 /// // NOTE: The try operator `?` is required for errors to be handled by proptest.
575 /// mock_service
576 /// .expect_request_that(|request| *request > 0).await?
577 /// .respond("OK");
578 ///
579 /// prop_assert!(matches!(call.await, Ok(Ok("OK"))));
580 /// #
581 /// # Ok::<(), TestCaseError>(())
582 /// # };
583 /// # test_code().await
584 /// # }).unwrap();
585 /// ```
586 pub async fn expect_request_that(
587 &mut self,
588 condition: impl FnOnce(&Request) -> bool,
589 ) -> Result<ResponseSender<Request, Response, Error>, TestCaseError>
590 where
591 Request: Debug,
592 {
593 let response_sender = self.next_request().await?;
594
595 prop_assert!(
596 condition(&response_sender.request),
597 "condition was false for request: {:?},\n \
598 in {}",
599 &response_sender.request,
600 std::any::type_name::<Self>(),
601 );
602
603 Ok(response_sender)
604 }
605
606 /// Expect no requests to be received.
607 ///
608 /// The internal queue of received requests should be empty, and no new requests should arrive
609 /// for the max delay time configured by [`MockServiceBuilder::with_max_request_delay`].
610 ///
611 /// If the queue is not empty or if a request is received before the max request delay timeout
612 /// expires, an error generated by a [`mod@proptest`] assertion is returned.
613 ///
614 /// # Example
615 ///
616 /// ```
617 /// # use proptest::prelude::TestCaseError;
618 /// # use tower::ServiceExt;
619 /// #
620 /// # use zebra_test::mock_service::MockService;
621 /// #
622 /// # let reactor = tokio::runtime::Builder::new_current_thread()
623 /// # .enable_all()
624 /// # .build()
625 /// # .expect("Failed to build Tokio runtime");
626 /// #
627 /// # reactor.block_on(async {
628 /// # let test_code = || async {
629 /// # let mut mock_service: MockService<(), (), _> =
630 /// # MockService::build().for_prop_tests();
631 /// #
632 /// // NOTE: The try operator `?` is required for errors to be handled by proptest.
633 /// mock_service.expect_no_requests().await?;
634 /// #
635 /// # Ok::<(), TestCaseError>(())
636 /// # };
637 /// # test_code().await
638 /// # }).unwrap();
639 /// ```
640 pub async fn expect_no_requests(&mut self) -> Result<(), TestCaseError>
641 where
642 Request: Debug,
643 {
644 match self.try_next_request().await {
645 Some(response_sender) => {
646 prop_assert!(
647 false,
648 "received an unexpected request: {:?},\n \
649 in {}",
650 response_sender.request,
651 std::any::type_name::<Self>(),
652 );
653 unreachable!("prop_assert!(false) returns an early error");
654 }
655 None => Ok(()),
656 }
657 }
658
659 /// A helper method to get the next request from the queue.
660 ///
661 /// Returns the next request in the internal queue or waits at most the max delay time
662 /// configured by [`MockServiceBuilder::with_max_request_delay`] for a new request to be
663 /// received, and then returns that.
664 ///
665 /// If the queue is empty and a request is not received before the max request delay timeout
666 /// expires, an error generated by a [`mod@proptest`] assertion is returned.
667 async fn next_request(
668 &mut self,
669 ) -> Result<ResponseSender<Request, Response, Error>, TestCaseError> {
670 match self.try_next_request().await {
671 Some(request) => Ok(request),
672 None => {
673 prop_assert!(
674 false,
675 "timeout while waiting for a request\n \
676 in {}",
677 std::any::type_name::<Self>(),
678 );
679 unreachable!("prop_assert!(false) returns an early error");
680 }
681 }
682 }
683
684 /// Returns a count of the number of times this service has been polled.
685 ///
686 /// Note: The poll count wraps around on overflow.
687 pub fn poll_count(&self) -> usize {
688 self.poll_count.load(Ordering::SeqCst)
689 }
690}
691
692/// Code that is independent of the assertions used in [`MockService`].
693impl<Request, Response, Assertion, Error> MockService<Request, Response, Assertion, Error> {
694 /// Try to get the next request received.
695 ///
696 /// Returns the next element in the queue. If the queue is empty, waits at most the max request
697 /// delay configured by [`MockServiceBuilder::with_max_request_delay`] for a request, and
698 /// returns it.
699 ///
700 /// If no request is received, returns `None`.
701 ///
702 /// If too many requests are received and the queue fills up, the oldest requests are dropped
703 /// and ignored. This means that calling this may not receive the next request if the queue is
704 /// not dimensioned properly with the [`MockServiceBuilder::with_proxy_channel_size`] method.
705 pub async fn try_next_request(&mut self) -> Option<ResponseSender<Request, Response, Error>> {
706 loop {
707 match timeout(self.max_request_delay, self.receiver.recv()).await {
708 Ok(Ok(item)) => {
709 if let Some(proxy_item) = item.lock().await.take() {
710 return Some(proxy_item);
711 }
712 }
713 Ok(Err(RecvError::Lagged(_))) => continue,
714 Ok(Err(RecvError::Closed)) => unreachable!("sender is never closed"),
715 Err(_timeout) => return None,
716 }
717 }
718 }
719}
720
721impl<Request, Response, Assertion, Error> Clone
722 for MockService<Request, Response, Assertion, Error>
723{
724 /// Clones the [`MockService`].
725 ///
726 /// This is a cheap operation, because it simply clones the [`broadcast`] channel endpoints.
727 fn clone(&self) -> Self {
728 MockService {
729 receiver: self.sender.subscribe(),
730 sender: self.sender.clone(),
731 poll_count: self.poll_count.clone(),
732 max_request_delay: self.max_request_delay,
733 _assertion_type: PhantomData,
734 }
735 }
736}
737
738impl<Request, Response, Error> ResponseSender<Request, Response, Error> {
739 /// Create a [`ResponseSender`] for a given `request`.
740 fn new(request: Request) -> (Self, oneshot::Receiver<Result<Response, Error>>) {
741 let (response_sender, response_receiver) = oneshot::channel();
742
743 (
744 ResponseSender {
745 request,
746 response_sender,
747 },
748 response_receiver,
749 )
750 }
751
752 /// Access the `request` that's awaiting a response.
753 pub fn request(&self) -> &Request {
754 &self.request
755 }
756
757 /// Respond to the request using a fixed response value.
758 ///
759 /// The `response` can be of the `Response` type or a [`Result`]. This allows sending an error
760 /// representing an error while processing the request.
761 ///
762 /// This method takes ownership of the [`ResponseSender`] so that only one response can be
763 /// sent.
764 ///
765 /// # Panics
766 ///
767 /// If one of the `respond*` methods isn't called, the [`MockService`] might panic with a
768 /// timeout error.
769 ///
770 /// # Example
771 ///
772 /// ```
773 /// # use zebra_test::mock_service::MockService;
774 /// # use tower::{Service, ServiceExt};
775 /// #
776 /// # #[derive(Debug, PartialEq, Eq)]
777 /// # struct Request;
778 /// #
779 /// # let reactor = tokio::runtime::Builder::new_current_thread()
780 /// # .enable_all()
781 /// # .build()
782 /// # .expect("Failed to build Tokio runtime");
783 /// #
784 /// # reactor.block_on(async {
785 /// // Mock a service with a `String` as the service `Error` type.
786 /// let mut mock_service: MockService<_, _, _, String> =
787 /// MockService::build().for_unit_tests();
788 ///
789 /// # let mut service = mock_service.clone();
790 /// # let task = tokio::spawn(async move {
791 /// # let first_call_result = (&mut service).oneshot(Request).await;
792 /// # let second_call_result = service.oneshot(Request).await;
793 /// #
794 /// # (first_call_result, second_call_result)
795 /// # });
796 /// #
797 /// mock_service
798 /// .expect_request(Request)
799 /// .await
800 /// .respond("Received Request".to_owned());
801 ///
802 /// mock_service
803 /// .expect_request(Request)
804 /// .await
805 /// .respond(Err("Duplicate request"));
806 /// # });
807 /// ```
808 pub fn respond(self, response: impl ResponseResult<Response, Error>) {
809 let _ = self.response_sender.send(response.into_result());
810 }
811
812 /// Respond to the request by calculating a value from the request.
813 ///
814 /// The response can be of the `Response` type or a [`Result`]. This allows sending an error
815 /// representing an error while processing the request.
816 ///
817 /// This method takes ownership of the [`ResponseSender`] so that only one response can be
818 /// sent.
819 ///
820 /// # Panics
821 ///
822 /// If one of the `respond*` methods isn't called, the [`MockService`] might panic with a
823 /// timeout error.
824 ///
825 /// # Example
826 ///
827 /// ```
828 /// # use zebra_test::mock_service::MockService;
829 /// # use tower::{Service, ServiceExt};
830 /// #
831 /// # #[derive(Debug, PartialEq, Eq)]
832 /// # struct Request;
833 /// #
834 /// # let reactor = tokio::runtime::Builder::new_current_thread()
835 /// # .enable_all()
836 /// # .build()
837 /// # .expect("Failed to build Tokio runtime");
838 /// #
839 /// # reactor.block_on(async {
840 /// // Mock a service with a `String` as the service `Error` type.
841 /// let mut mock_service: MockService<_, _, _, String> =
842 /// MockService::build().for_unit_tests();
843 ///
844 /// # let mut service = mock_service.clone();
845 /// # let task = tokio::spawn(async move {
846 /// # let first_call_result = (&mut service).oneshot(Request).await;
847 /// # let second_call_result = service.oneshot(Request).await;
848 /// #
849 /// # (first_call_result, second_call_result)
850 /// # });
851 /// #
852 /// mock_service
853 /// .expect_request(Request)
854 /// .await
855 /// .respond_with(|req| format!("Received: {req:?}"));
856 ///
857 /// mock_service
858 /// .expect_request(Request)
859 /// .await
860 /// .respond_with(|req| Err(format!("Duplicate request: {req:?}")));
861 /// # });
862 /// ```
863 pub fn respond_with<F, R>(self, response_fn: F)
864 where
865 F: FnOnce(&Request) -> R,
866 R: ResponseResult<Response, Error>,
867 {
868 let response_result = response_fn(self.request()).into_result();
869 let _ = self.response_sender.send(response_result);
870 }
871
872 /// Respond to the request using a fixed error value.
873 ///
874 /// The `error` must be the `Error` type. This helps avoid type resolution issues in the
875 /// compiler.
876 ///
877 /// This method takes ownership of the [`ResponseSender`] so that only one response can be
878 /// sent.
879 ///
880 /// # Panics
881 ///
882 /// If one of the `respond*` methods isn't called, the [`MockService`] might panic with a
883 /// timeout error.
884 ///
885 /// # Example
886 ///
887 /// ```
888 /// # use zebra_test::mock_service::MockService;
889 /// # use tower::{Service, ServiceExt};
890 /// #
891 /// # #[derive(Debug, PartialEq, Eq)]
892 /// # struct Request;
893 /// # struct Response;
894 /// #
895 /// # let reactor = tokio::runtime::Builder::new_current_thread()
896 /// # .enable_all()
897 /// # .build()
898 /// # .expect("Failed to build Tokio runtime");
899 /// #
900 /// # reactor.block_on(async {
901 /// // Mock a service with a `String` as the service `Error` type.
902 /// let mut mock_service: MockService<Request, Response, _, String> =
903 /// MockService::build().for_unit_tests();
904 ///
905 /// # let mut service = mock_service.clone();
906 /// # let task = tokio::spawn(async move {
907 /// # let first_call_result = (&mut service).oneshot(Request).await;
908 /// # let second_call_result = service.oneshot(Request).await;
909 /// #
910 /// # (first_call_result, second_call_result)
911 /// # });
912 /// #
913 /// mock_service
914 /// .expect_request(Request)
915 /// .await
916 /// .respond_error("Duplicate request".to_string());
917 /// # });
918 /// ```
919 pub fn respond_error(self, error: Error) {
920 // TODO: impl ResponseResult for BoxError/Error trait when overlapping impls are
921 // better supported by the compiler
922 let _ = self.response_sender.send(Err(error));
923 }
924
925 /// Respond to the request by calculating an error from the request.
926 ///
927 /// The `error` must be the `Error` type. This helps avoid type resolution issues in the
928 /// compiler.
929 ///
930 /// This method takes ownership of the [`ResponseSender`] so that only one response can be
931 /// sent.
932 ///
933 /// # Panics
934 ///
935 /// If one of the `respond*` methods isn't called, the [`MockService`] might panic with a
936 /// timeout error.
937 ///
938 /// # Example
939 ///
940 /// ```
941 /// # use zebra_test::mock_service::MockService;
942 /// # use tower::{Service, ServiceExt};
943 /// #
944 /// # #[derive(Debug, PartialEq, Eq)]
945 /// # struct Request;
946 /// # struct Response;
947 /// #
948 /// # let reactor = tokio::runtime::Builder::new_current_thread()
949 /// # .enable_all()
950 /// # .build()
951 /// # .expect("Failed to build Tokio runtime");
952 /// #
953 /// # reactor.block_on(async {
954 /// // Mock a service with a `String` as the service `Error` type.
955 /// let mut mock_service: MockService<Request, Response, _, String> =
956 /// MockService::build().for_unit_tests();
957 ///
958 /// # let mut service = mock_service.clone();
959 /// # let task = tokio::spawn(async move {
960 /// # let first_call_result = (&mut service).oneshot(Request).await;
961 /// # let second_call_result = service.oneshot(Request).await;
962 /// #
963 /// # (first_call_result, second_call_result)
964 /// # });
965 /// #
966 /// mock_service
967 /// .expect_request(Request)
968 /// .await
969 /// .respond_with_error(|req| format!("Duplicate request: {req:?}"));
970 /// # });
971 /// ```
972 pub fn respond_with_error<F>(self, response_fn: F)
973 where
974 F: FnOnce(&Request) -> Error,
975 {
976 // TODO: impl ResponseResult for BoxError/Error trait when overlapping impls are
977 // better supported by the compiler
978 let response_result = Err(response_fn(self.request()));
979 let _ = self.response_sender.send(response_result);
980 }
981}
982
983/// A representation of an assertion type.
984///
985/// This trait is used to group the types of assertions that the [`MockService`] can do. There are
986/// currently two types that are used as type-system tags on the [`MockService`]:
987///
988/// - [`PanicAssertion`]
989/// - [`PropTestAssertion`]
990#[allow(dead_code)]
991trait AssertionType {}
992
993/// Represents normal Rust assertions that panic, like [`assert_eq`].
994pub enum PanicAssertion {}
995
996/// Represents [`mod@proptest`] assertions that return errors, like [`prop_assert_eq`].
997pub enum PropTestAssertion {}
998
999impl AssertionType for PanicAssertion {}
1000
1001impl AssertionType for PropTestAssertion {}
1002
1003/// A helper trait to improve ergonomics when sending a response.
1004///
1005/// This allows the [`ResponseSender::respond`] method to receive either a [`Result`] or just the
1006/// response type, which it automatically wraps in an `Ok` variant.
1007pub trait ResponseResult<Response, Error> {
1008 /// Converts the type into a [`Result`] that can be sent as a response.
1009 fn into_result(self) -> Result<Response, Error>;
1010}
1011
1012impl<Response, Error> ResponseResult<Response, Error> for Response {
1013 fn into_result(self) -> Result<Response, Error> {
1014 Ok(self)
1015 }
1016}
1017
1018impl<Response, SourceError, TargetError> ResponseResult<Response, TargetError>
1019 for Result<Response, SourceError>
1020where
1021 SourceError: Into<TargetError>,
1022{
1023 fn into_result(self) -> Result<Response, TargetError> {
1024 self.map_err(|source_error| source_error.into())
1025 }
1026}