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}