1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

#![allow(missing_docs)]
use std::time::Duration;

#[macro_export]
macro_rules! wait_for_any {
    ($wait:expr, $test:expr, $check:expr, $assert:expr) => {{
        loop {
            let o = $test;
            if !$wait.wait_any().await || $check(&o) {
                $assert(o);
                break;
            }
        }
    }};
}

#[macro_export]
macro_rules! wait_for_any_10s {
    ($test:expr, $check:expr, $assert:expr) => {
        let mut wait_for = $crate::test_utils::WaitForAny::ten_s();
        $crate::wait_for_any!(wait_for, $test, $check, $assert)
    };
}

#[macro_export]
macro_rules! wait_for_any_1m {
    ($test:expr, $check:expr, $assert:expr) => {
        let mut wait_for = $crate::test_utils::WaitForAny::one_m();
        $crate::wait_for_any!(wait_for, $test, $check, $assert)
    };
}

#[macro_export]
macro_rules! assert_retry {
    ($wait:expr, $test:expr, $check:expr $(, $reason:literal)?) => {
        $crate::wait_for_any!($wait, $test, $check, |x| assert!(x $(, $reason)?))
    };
}

#[macro_export]
macro_rules! assert_eq_retry {
    ($wait:expr, $test:expr, $check:expr $(, $reason:literal)?) => {
        $crate::wait_for_any!($wait, $test, |x| x == &$check, |x| assert_eq!(x, $check $(, $reason)?))
    };
}

#[macro_export]
macro_rules! assert_retry_10s {
    ($test:expr, $check:expr $(, $reason:literal)? $(,)?) => {
        let mut wait_for = $crate::test_utils::WaitForAny::ten_s();
        $crate::assert_retry!(wait_for, $test, $check  $(, $reason:literal)?)
    };
}

#[macro_export]
macro_rules! assert_eq_retry_10s {
    ($test:expr, $check:expr $(, $reason:literal)? $(,)?) => {
        let mut wait_for = $crate::test_utils::WaitForAny::ten_s();
        $crate::assert_eq_retry!(wait_for, $test, $check  $(, $reason:literal)?)
    };
}

#[macro_export]
macro_rules! assert_retry_1m {
    ($test:expr, $check:expr $(, $reason:literal)? $(,)?) => {
        let mut wait_for = $crate::test_utils::WaitForAny::one_m();
        $crate::assert_retry!(wait_for, $test, $check  $(, $reason:literal)?)
    };
}

#[macro_export]
macro_rules! assert_eq_retry_1m {
    ($test:expr, $check:expr $(, $reason:literal)? $(,)?) => {
        let mut wait_for = $crate::test_utils::WaitForAny::one_m();
        $crate::assert_eq_retry!(wait_for, $test, $check  $(, $reason:literal)?)
    };
}

#[macro_export]
macro_rules! assert_eq_retry_5m {
    ($test:expr, $check:expr $(, $reason:literal)? $(,)?) => {
        let mut wait_for = $crate::test_utils::WaitForAny::five_m();
        $crate::assert_eq_retry!(wait_for, $test, $check  $(, $reason:literal)?)
    };
}

#[derive(Debug, Clone)]
/// Generic waiting for some test property to
/// be true. This allows early exit from waiting when
/// the condition becomes true but will wait up to a
/// maximum if the condition is not true.
pub struct WaitForAny {
    num_attempts: usize,
    attempt: usize,
    delay: Duration,
}

impl WaitForAny {
    /// Create a new wait for from a number of attempts and delay in between attempts
    pub fn new(num_attempts: usize, delay: Duration) -> Self {
        Self {
            num_attempts,
            attempt: 0,
            delay,
        }
    }

    /// Wait for 10s checking every 100ms.
    pub fn ten_s() -> Self {
        const DELAY_PER_ATTEMPT: std::time::Duration = std::time::Duration::from_millis(100);
        Self::new(100, DELAY_PER_ATTEMPT)
    }

    /// Wait for 1 minute checking every 500ms.
    pub fn one_m() -> Self {
        const DELAY_PER_ATTEMPT: std::time::Duration = std::time::Duration::from_millis(500);
        Self::new(120, DELAY_PER_ATTEMPT)
    }

    /// Wait for 5 minutes checking every 1000ms.
    pub fn five_m() -> Self {
        const DELAY_PER_ATTEMPT: std::time::Duration = std::time::Duration::from_millis(1000);
        Self::new(60 * 5, DELAY_PER_ATTEMPT)
    }

    /// Wait for some time before trying again.
    /// Will return false when you should stop waiting.
    #[tracing::instrument(skip(self))]
    pub async fn wait_any(&mut self) -> bool {
        if self.attempt >= self.num_attempts {
            return false;
        }
        self.attempt += 1;
        tracing::debug!(attempt = ?self.attempt, out_of = ?self.num_attempts, delaying_for = ?self.delay);
        tokio::time::sleep(self.delay).await;
        true
    }
}