async_time_mock_tokio/
timeout.rs

1use crate::elapsed::Elapsed;
2use pin_project::pin_project;
3use std::fmt::{Debug, Formatter};
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8#[pin_project(project = ProjectedTimeout)]
9pub enum Timeout<T> {
10	Real(#[pin] tokio::time::Timeout<T>),
11	#[cfg(feature = "mock")]
12	Mock(#[pin] async_time_mock_core::Timeout<T>),
13}
14
15impl<T> Timeout<T> {
16	pub fn get_ref(&self) -> &T {
17		use Timeout::*;
18		match self {
19			Real(timeout) => timeout.get_ref(),
20			#[cfg(feature = "mock")]
21			Mock(timeout) => timeout.get_ref(),
22		}
23	}
24
25	pub fn get_mut(&mut self) -> &mut T {
26		use Timeout::*;
27		match self {
28			Real(timeout) => timeout.get_mut(),
29			#[cfg(feature = "mock")]
30			Mock(timeout) => timeout.get_mut(),
31		}
32	}
33
34	pub fn into_inner(self) -> T {
35		use Timeout::*;
36		match self {
37			Real(timeout) => timeout.into_inner(),
38			#[cfg(feature = "mock")]
39			Mock(timeout) => timeout.into_inner(),
40		}
41	}
42}
43
44impl<T: Debug> Debug for Timeout<T> {
45	fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
46		use Timeout::*;
47		match self {
48			Real(timeout) => Debug::fmt(timeout, formatter),
49			#[cfg(feature = "mock")]
50			Mock(_) => formatter.debug_struct("Mock(Timeout)").finish(),
51		}
52	}
53}
54
55impl<T> Future for Timeout<T>
56where
57	T: Future,
58{
59	type Output = Result<T::Output, Elapsed>;
60
61	fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
62		use ProjectedTimeout::*;
63		match self.project() {
64			Real(timeout) => timeout.poll(context).map(|result| result.map_err(Into::into)),
65			#[cfg(feature = "mock")]
66			Mock(timeout) => timeout.poll(context).map(|result| result.map_err(Into::into)),
67		}
68	}
69}
70
71impl<T> From<tokio::time::Timeout<T>> for Timeout<T> {
72	fn from(timeout: tokio::time::Timeout<T>) -> Self {
73		Self::Real(timeout)
74	}
75}
76
77#[cfg(feature = "mock")]
78impl<T> From<async_time_mock_core::Timeout<T>> for Timeout<T> {
79	fn from(timeout: async_time_mock_core::Timeout<T>) -> Self {
80		Self::Mock(timeout)
81	}
82}