Skip to main content

vector/test_util/
mod.rs

1#![allow(missing_docs)]
2
3use std::{
4    collections::HashMap,
5    convert::Infallible,
6    fs::File,
7    future::{Future, ready},
8    io::Read,
9    iter,
10    net::SocketAddr,
11    path::{Path, PathBuf},
12    pin::Pin,
13    sync::{
14        Arc,
15        atomic::{AtomicUsize, Ordering},
16    },
17    task::{Context, Poll, ready},
18};
19
20use chrono::{DateTime, SubsecRound, Utc};
21use futures::{FutureExt, SinkExt, Stream, StreamExt, TryStreamExt, stream, task::noop_waker_ref};
22use openssl::ssl::{SslConnector, SslFiletype, SslMethod, SslVerifyMode};
23use rand::{Rng, rng};
24use rand_distr::Alphanumeric;
25use tokio::{
26    io::{AsyncRead, AsyncWrite, AsyncWriteExt, Result as IoResult},
27    net::{TcpListener, TcpStream, ToSocketAddrs},
28    runtime,
29    sync::oneshot,
30    task::JoinHandle,
31    time::{Duration, Instant, sleep},
32};
33use tokio_stream::wrappers::TcpListenerStream;
34#[cfg(unix)]
35use tokio_stream::wrappers::UnixListenerStream;
36use tokio_util::codec::{Encoder, FramedRead, FramedWrite, LinesCodec};
37use vector_common::decompression::CappedDecoder;
38use vector_lib::{
39    buffers::topology::channel::LimitedReceiver,
40    event::{
41        BatchNotifier, BatchStatusReceiver, Event, EventArray, LogEvent, Metric, MetricKind,
42        MetricTags, MetricValue,
43    },
44};
45
46use crate::{
47    config::{Config, GenerateConfig},
48    topology::{RunningTopology, ShutdownErrorReceiver},
49    trace,
50};
51
52const WAIT_FOR_SECS: u64 = 5; // The default time to wait in `wait_for`
53const WAIT_FOR_MIN_MILLIS: u64 = 5; // The minimum time to pause before retrying
54const WAIT_FOR_MAX_MILLIS: u64 = 500; // The maximum time to pause before retrying
55
56pub mod addr;
57pub mod compression;
58pub mod stats;
59
60#[cfg(any(test, feature = "test-utils"))]
61pub mod components;
62#[cfg(test)]
63pub mod http;
64#[cfg(test)]
65pub mod integration;
66#[cfg(test)]
67pub mod metrics;
68#[cfg(test)]
69pub mod mock;
70
71#[macro_export]
72macro_rules! assert_downcast_matches {
73    ($e:expr_2021, $t:ty, $v:pat) => {{
74        match $e.downcast_ref::<$t>() {
75            Some($v) => (),
76            got => panic!("Assertion failed: got wrong error variant {:?}", got),
77        }
78    }};
79}
80
81#[macro_export]
82macro_rules! log_event {
83    ($($key:expr_2021 => $value:expr_2021),*  $(,)?) => {
84        #[allow(unused_variables)]
85        {
86            let mut event = $crate::event::Event::Log($crate::event::LogEvent::default());
87            let log = event.as_mut_log();
88            $(
89                log.insert(::vrl::event_path!($key), $value);
90            )*
91            event
92        }
93    };
94}
95
96pub fn test_generate_config<T>()
97where
98    for<'de> T: GenerateConfig + serde::Deserialize<'de>,
99{
100    let cfg = toml::to_string(&T::generate_config()).unwrap();
101
102    toml::from_str::<T>(&cfg)
103        .unwrap_or_else(|e| panic!("Invalid config generated from string:\n\n{e}\n'{cfg}'"));
104}
105
106pub fn open_fixture(path: impl AsRef<Path>) -> crate::Result<serde_json::Value> {
107    let test_file = File::open(path)?;
108    let value: serde_json::Value = serde_json::from_reader(test_file)?;
109    Ok(value)
110}
111
112pub fn trace_init() {
113    #[cfg(unix)]
114    let color = {
115        use std::io::IsTerminal;
116        std::io::stdout().is_terminal()
117            || std::env::var("NEXTEST")
118                .ok()
119                .and(Some(true))
120                .unwrap_or(false)
121    };
122    // Windows: ANSI colors are not supported by cmd.exe
123    // Color is false for everything except unix.
124    #[cfg(not(unix))]
125    let color = false;
126
127    let levels = std::env::var("VECTOR_LOG").unwrap_or_else(|_| "error".to_string());
128
129    trace::init(color, false, &levels, 10, None);
130
131    // Initialize metrics as well
132    vector_lib::metrics::init_test();
133}
134
135pub async fn send_lines(
136    addr: SocketAddr,
137    lines: impl IntoIterator<Item = String>,
138) -> Result<SocketAddr, Infallible> {
139    send_encodable(addr, LinesCodec::new(), lines).await
140}
141
142pub async fn send_encodable<I, E: From<std::io::Error> + std::fmt::Debug>(
143    addr: SocketAddr,
144    encoder: impl Encoder<I, Error = E>,
145    lines: impl IntoIterator<Item = I>,
146) -> Result<SocketAddr, Infallible> {
147    let stream = TcpStream::connect(&addr).await.unwrap();
148
149    let local_addr = stream.local_addr().unwrap();
150
151    let mut sink = FramedWrite::new(stream, encoder);
152
153    let mut lines = stream::iter(lines).map(Ok);
154    sink.send_all(&mut lines).await.unwrap();
155
156    let stream = sink.get_mut();
157    stream.shutdown().await.unwrap();
158
159    Ok(local_addr)
160}
161
162pub async fn send_lines_tls(
163    addr: SocketAddr,
164    host: String,
165    lines: impl Iterator<Item = String>,
166    ca: impl Into<Option<&Path>>,
167    client_cert: impl Into<Option<&Path>>,
168    client_key: impl Into<Option<&Path>>,
169) -> Result<SocketAddr, Infallible> {
170    let stream = TcpStream::connect(&addr).await.unwrap();
171
172    let local_addr = stream.local_addr().unwrap();
173
174    let mut connector = SslConnector::builder(SslMethod::tls()).unwrap();
175    if let Some(ca) = ca.into() {
176        connector.set_ca_file(ca).unwrap();
177    } else {
178        connector.set_verify(SslVerifyMode::NONE);
179    }
180
181    if let Some(cert_file) = client_cert.into() {
182        connector.set_certificate_chain_file(cert_file).unwrap();
183    }
184
185    if let Some(key_file) = client_key.into() {
186        connector
187            .set_private_key_file(key_file, SslFiletype::PEM)
188            .unwrap();
189    }
190
191    let ssl = connector
192        .build()
193        .configure()
194        .unwrap()
195        .into_ssl(&host)
196        .unwrap();
197
198    let mut stream = tokio_openssl::SslStream::new(ssl, stream).unwrap();
199    Pin::new(&mut stream).connect().await.unwrap();
200    let mut sink = FramedWrite::new(stream, LinesCodec::new());
201
202    let mut lines = stream::iter(lines).map(Ok);
203    sink.send_all(&mut lines).await.unwrap();
204
205    let stream = sink.get_mut().get_mut();
206    stream.shutdown().await.unwrap();
207
208    Ok(local_addr)
209}
210
211pub fn temp_file() -> PathBuf {
212    let path = std::env::temp_dir();
213    let file_name = random_string(16);
214    path.join(file_name + ".log")
215}
216
217pub fn temp_dir() -> PathBuf {
218    let path = std::env::temp_dir();
219    let dir_name = random_string(16);
220    path.join(dir_name)
221}
222
223pub fn random_table_name() -> String {
224    format!("test_{}", random_string(10).to_lowercase())
225}
226
227pub fn map_event_batch_stream(
228    stream: impl Stream<Item = Event>,
229    batch: Option<BatchNotifier>,
230) -> impl Stream<Item = EventArray> {
231    stream.map(move |event| event.with_batch_notifier_option(&batch).into())
232}
233
234// TODO refactor to have a single implementation for `Event`, `LogEvent` and `Metric`.
235fn map_batch_stream(
236    stream: impl Stream<Item = LogEvent>,
237    batch: Option<BatchNotifier>,
238) -> impl Stream<Item = EventArray> {
239    stream.map(move |log| vec![log.with_batch_notifier_option(&batch)].into())
240}
241
242pub fn generate_lines_with_stream<Gen: FnMut(usize) -> String>(
243    generator: Gen,
244    count: usize,
245    batch: Option<BatchNotifier>,
246) -> (Vec<String>, impl Stream<Item = EventArray>) {
247    let lines = (0..count).map(generator).collect::<Vec<_>>();
248    let stream = map_batch_stream(
249        stream::iter(lines.clone()).map(LogEvent::from_str_legacy),
250        batch,
251    );
252    (lines, stream)
253}
254
255pub fn random_lines_with_stream(
256    len: usize,
257    count: usize,
258    batch: Option<BatchNotifier>,
259) -> (Vec<String>, impl Stream<Item = EventArray>) {
260    let generator = move |_| random_string(len);
261    generate_lines_with_stream(generator, count, batch)
262}
263
264pub fn generate_events_with_stream<Gen: FnMut(usize) -> Event>(
265    generator: Gen,
266    count: usize,
267    batch: Option<BatchNotifier>,
268) -> (Vec<Event>, impl Stream<Item = EventArray>) {
269    let events = (0..count).map(generator).collect::<Vec<_>>();
270    let stream = map_batch_stream(
271        stream::iter(events.clone()).map(|event| event.into_log()),
272        batch,
273    );
274    (events, stream)
275}
276
277pub fn random_metrics_with_stream(
278    count: usize,
279    batch: Option<BatchNotifier>,
280    tags: Option<MetricTags>,
281) -> (Vec<Event>, impl Stream<Item = EventArray>) {
282    random_metrics_with_stream_timestamp(
283        count,
284        batch,
285        tags,
286        Utc::now().trunc_subsecs(3),
287        std::time::Duration::from_secs(2),
288    )
289}
290
291/// Generates event metrics with the provided tags and timestamp.
292///
293/// # Parameters
294/// - `count`: the number of metrics to generate
295/// - `batch`: the batch notifier to use with the stream
296/// - `tags`: the tags to apply to each metric event
297/// - `timestamp`: the timestamp to use for each metric event
298/// - `timestamp_offset`: the offset from the `timestamp` to use for each additional metric
299///
300/// # Returns
301/// A tuple of the generated metric events and the stream of the generated events
302pub fn random_metrics_with_stream_timestamp(
303    count: usize,
304    batch: Option<BatchNotifier>,
305    tags: Option<MetricTags>,
306    timestamp: DateTime<Utc>,
307    timestamp_offset: std::time::Duration,
308) -> (Vec<Event>, impl Stream<Item = EventArray>) {
309    let events: Vec<_> = (0..count)
310        .map(|index| {
311            let ts = timestamp + (timestamp_offset * index as u32);
312            Event::Metric(
313                Metric::new(
314                    format!("counter_{}", rng().random::<u32>()),
315                    MetricKind::Incremental,
316                    MetricValue::Counter {
317                        value: index as f64,
318                    },
319                )
320                .with_timestamp(Some(ts))
321                .with_tags(tags.clone()),
322            )
323            // this ensures we get Origin Metadata, with an undefined service but that's ok.
324            .with_source_type("a_source_like_none_other")
325        })
326        .collect();
327
328    let stream = map_event_batch_stream(stream::iter(events.clone()), batch);
329    (events, stream)
330}
331
332pub fn random_events_with_stream(
333    len: usize,
334    count: usize,
335    batch: Option<BatchNotifier>,
336) -> (Vec<Event>, impl Stream<Item = EventArray>) {
337    let events = (0..count)
338        .map(|_| Event::from(LogEvent::from_str_legacy(random_string(len))))
339        .collect::<Vec<_>>();
340    let stream = map_batch_stream(
341        stream::iter(events.clone()).map(|event| event.into_log()),
342        batch,
343    );
344    (events, stream)
345}
346
347pub fn random_updated_events_with_stream<F>(
348    len: usize,
349    count: usize,
350    batch: Option<BatchNotifier>,
351    update_fn: F,
352) -> (Vec<Event>, impl Stream<Item = EventArray>)
353where
354    F: Fn((usize, LogEvent)) -> LogEvent,
355{
356    let events = (0..count)
357        .map(|_| LogEvent::from_str_legacy(random_string(len)))
358        .enumerate()
359        .map(update_fn)
360        .map(Event::Log)
361        .collect::<Vec<_>>();
362    let stream = map_batch_stream(
363        stream::iter(events.clone()).map(|event| event.into_log()),
364        batch,
365    );
366    (events, stream)
367}
368
369pub fn create_events_batch_with_fn<F: Fn() -> Event>(
370    create_event_fn: F,
371    num_events: usize,
372) -> (Vec<Event>, BatchStatusReceiver) {
373    let mut events = (0..num_events)
374        .map(|_| create_event_fn())
375        .collect::<Vec<_>>();
376    let receiver = BatchNotifier::apply_to(&mut events);
377    (events, receiver)
378}
379
380pub fn random_string(len: usize) -> String {
381    rng()
382        .sample_iter(&Alphanumeric)
383        .take(len)
384        .map(char::from)
385        .collect::<String>()
386}
387
388pub fn random_lines(len: usize) -> impl Iterator<Item = String> {
389    iter::repeat_with(move || random_string(len))
390}
391
392pub fn random_map(max_size: usize, field_len: usize) -> HashMap<String, String> {
393    let size = rng().random_range(0..max_size);
394
395    (0..size)
396        .map(move |_| (random_string(field_len), random_string(field_len)))
397        .collect()
398}
399
400pub fn random_maps(
401    max_size: usize,
402    field_len: usize,
403) -> impl Iterator<Item = HashMap<String, String>> {
404    iter::repeat_with(move || random_map(max_size, field_len))
405}
406
407pub async fn collect_n<S>(rx: S, n: usize) -> Vec<S::Item>
408where
409    S: Stream,
410{
411    rx.take(n).collect().await
412}
413
414pub async fn collect_n_stream<T, S: Stream<Item = T> + Unpin>(stream: &mut S, n: usize) -> Vec<T> {
415    let mut events = Vec::with_capacity(n);
416
417    while events.len() < n {
418        let e = stream.next().await.unwrap();
419        events.push(e);
420    }
421    events
422}
423
424pub async fn collect_ready<S>(mut rx: S) -> Vec<S::Item>
425where
426    S: Stream + Unpin,
427{
428    let waker = noop_waker_ref();
429    let mut cx = Context::from_waker(waker);
430
431    let mut vec = Vec::new();
432    loop {
433        match rx.poll_next_unpin(&mut cx) {
434            Poll::Ready(Some(item)) => vec.push(item),
435            Poll::Ready(None) | Poll::Pending => return vec,
436        }
437    }
438}
439
440pub async fn collect_limited<T: Send + 'static>(mut rx: LimitedReceiver<T>) -> Vec<T> {
441    let mut items = Vec::new();
442    while let Some(item) = rx.next().await {
443        items.push(item);
444    }
445    items
446}
447
448pub async fn collect_n_limited<T: Send + 'static>(mut rx: LimitedReceiver<T>, n: usize) -> Vec<T> {
449    let mut items = Vec::new();
450    while items.len() < n {
451        match rx.next().await {
452            Some(item) => items.push(item),
453            None => break,
454        }
455    }
456    items
457}
458
459pub fn lines_from_file<P: AsRef<Path>>(path: P) -> Vec<String> {
460    trace!(message = "Reading file.", path = %path.as_ref().display());
461    let mut file = File::open(path).unwrap();
462    let mut output = String::new();
463    file.read_to_string(&mut output).unwrap();
464    output.lines().map(|s| s.to_owned()).collect()
465}
466
467pub fn lines_from_gzip_file<P: AsRef<Path>>(path: P) -> Vec<String> {
468    trace!(message = "Reading gzip file.", path = %path.as_ref().display());
469    let mut file = File::open(path).unwrap();
470    let mut gzip_bytes = Vec::new();
471    file.read_to_end(&mut gzip_bytes).unwrap();
472    let output = CappedDecoder::gzip(&gzip_bytes[..]).decompress().unwrap();
473    String::from_utf8(output)
474        .unwrap()
475        .lines()
476        .map(|s| s.to_owned())
477        .collect()
478}
479
480#[cfg(test)]
481pub fn lines_from_zstd_file<P: AsRef<Path>>(path: P) -> Vec<String> {
482    trace!(message = "Reading zstd file.", path = %path.as_ref().display());
483    let file = File::open(path).unwrap();
484    let output = CappedDecoder::zstd(file).unwrap().decompress().unwrap();
485    String::from_utf8(output)
486        .unwrap()
487        .lines()
488        .map(|s| s.to_owned())
489        .collect()
490}
491
492pub fn runtime() -> runtime::Runtime {
493    runtime::Builder::new_multi_thread()
494        .enable_all()
495        .build()
496        .unwrap()
497}
498
499// Wait for a Future to resolve, or the duration to elapse (will panic)
500pub async fn wait_for_duration<F, Fut>(mut f: F, duration: Duration)
501where
502    F: FnMut() -> Fut,
503    Fut: Future<Output = bool> + Send + 'static,
504{
505    let started = Instant::now();
506    let mut delay = WAIT_FOR_MIN_MILLIS;
507    while !f().await {
508        sleep(Duration::from_millis(delay)).await;
509        if started.elapsed() > duration {
510            panic!("Timed out while waiting");
511        }
512        // quadratic backoff up to a maximum delay
513        delay = (delay * 2).min(WAIT_FOR_MAX_MILLIS);
514    }
515}
516
517// Wait for 5 seconds
518pub async fn wait_for<F, Fut>(f: F)
519where
520    F: FnMut() -> Fut,
521    Fut: Future<Output = bool> + Send + 'static,
522{
523    wait_for_duration(f, Duration::from_secs(WAIT_FOR_SECS)).await
524}
525
526// Wait (for 5 secs) for a TCP socket to be reachable
527pub async fn wait_for_tcp<A>(addr: A)
528where
529    A: ToSocketAddrs + Clone + Send + 'static,
530{
531    wait_for(move || {
532        let addr = addr.clone();
533        async move { TcpStream::connect(addr).await.is_ok() }
534    })
535    .await
536}
537
538// Allows specifying a custom duration to wait for a TCP socket to be reachable
539pub async fn wait_for_tcp_duration(addr: SocketAddr, duration: Duration) {
540    wait_for_duration(
541        || async move { TcpStream::connect(addr).await.is_ok() },
542        duration,
543    )
544    .await
545}
546
547pub async fn wait_for_atomic_usize<T, F>(value: T, unblock: F)
548where
549    T: AsRef<AtomicUsize>,
550    F: Fn(usize) -> bool,
551{
552    let value = value.as_ref();
553    wait_for(|| ready(unblock(value.load(Ordering::SeqCst)))).await
554}
555
556pub async fn wait_for_atomic_usize_timeout_ms<T, F>(value: T, unblock: F, timeout_ms: u64)
557where
558    T: AsRef<AtomicUsize>,
559    F: Fn(usize) -> bool,
560{
561    let value = value.as_ref();
562    wait_for_duration(
563        || ready(unblock(value.load(Ordering::SeqCst))),
564        Duration::from_millis(timeout_ms),
565    )
566    .await
567}
568
569// Retries a func every `retry` duration until given an Ok(T); panics after `until` elapses
570pub async fn retry_until<'a, F, Fut, T, E>(mut f: F, retry: Duration, until: Duration) -> T
571where
572    F: FnMut() -> Fut,
573    Fut: Future<Output = Result<T, E>> + Send + 'a,
574{
575    let started = Instant::now();
576    while started.elapsed() < until {
577        match f().await {
578            Ok(res) => return res,
579            Err(_) => tokio::time::sleep(retry).await,
580        }
581    }
582    panic!("Timeout")
583}
584
585pub struct CountReceiver<T> {
586    count: Arc<AtomicUsize>,
587    trigger: Option<oneshot::Sender<()>>,
588    connected: Option<oneshot::Receiver<()>>,
589    handle: JoinHandle<Vec<T>>,
590}
591
592impl<T: Send + 'static> CountReceiver<T> {
593    pub fn count(&self) -> usize {
594        self.count.load(Ordering::Relaxed)
595    }
596
597    /// Succeeds once first connection has been made.
598    pub async fn connected(&mut self) {
599        if let Some(tripwire) = self.connected.take() {
600            tripwire.await.unwrap();
601        }
602    }
603
604    fn new<F, Fut>(make_fut: F) -> CountReceiver<T>
605    where
606        F: FnOnce(Arc<AtomicUsize>, oneshot::Receiver<()>, oneshot::Sender<()>) -> Fut,
607        Fut: Future<Output = Vec<T>> + Send + 'static,
608    {
609        let count = Arc::new(AtomicUsize::new(0));
610        let (trigger, tripwire) = oneshot::channel();
611        let (trigger_connected, connected) = oneshot::channel();
612
613        CountReceiver {
614            count: Arc::clone(&count),
615            trigger: Some(trigger),
616            connected: Some(connected),
617            handle: tokio::spawn(make_fut(count, tripwire, trigger_connected)),
618        }
619    }
620
621    pub fn receive_items_stream<S, F, Fut>(make_stream: F) -> CountReceiver<T>
622    where
623        S: Stream<Item = T> + Send + 'static,
624        F: FnOnce(oneshot::Receiver<()>, oneshot::Sender<()>) -> Fut + Send + 'static,
625        Fut: Future<Output = S> + Send + 'static,
626    {
627        CountReceiver::new(|count, tripwire, connected| async move {
628            let stream = make_stream(tripwire, connected).await;
629            stream
630                .inspect(move |_| {
631                    count.fetch_add(1, Ordering::Relaxed);
632                })
633                .collect::<Vec<T>>()
634                .await
635        })
636    }
637}
638
639impl<T> Future for CountReceiver<T> {
640    type Output = Vec<T>;
641
642    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
643        let this = self.get_mut();
644        if let Some(trigger) = this.trigger.take() {
645            _ = trigger.send(());
646        }
647
648        let result = ready!(this.handle.poll_unpin(cx));
649        Poll::Ready(result.unwrap())
650    }
651}
652
653impl CountReceiver<String> {
654    pub fn receive_lines(addr: SocketAddr) -> CountReceiver<String> {
655        CountReceiver::new(|count, tripwire, connected| async move {
656            let listener = TcpListener::bind(addr).await.unwrap();
657            CountReceiver::receive_lines_stream(
658                TcpListenerStream::new(listener),
659                count,
660                tripwire,
661                Some(connected),
662            )
663            .await
664        })
665    }
666
667    #[cfg(unix)]
668    pub fn receive_lines_unix<P>(path: P) -> CountReceiver<String>
669    where
670        P: AsRef<Path> + Send + 'static,
671    {
672        CountReceiver::new(|count, tripwire, connected| async move {
673            let listener = tokio::net::UnixListener::bind(path).unwrap();
674            CountReceiver::receive_lines_stream(
675                UnixListenerStream::new(listener),
676                count,
677                tripwire,
678                Some(connected),
679            )
680            .await
681        })
682    }
683
684    async fn receive_lines_stream<S, T>(
685        stream: S,
686        count: Arc<AtomicUsize>,
687        tripwire: oneshot::Receiver<()>,
688        mut connected: Option<oneshot::Sender<()>>,
689    ) -> Vec<String>
690    where
691        S: Stream<Item = IoResult<T>>,
692        T: AsyncWrite + AsyncRead,
693    {
694        stream
695            .take_until(tripwire)
696            .map_ok(|socket| FramedRead::new(socket, LinesCodec::new()))
697            .map(|x| {
698                connected.take().map(|trigger| trigger.send(()));
699                x.unwrap()
700            })
701            .flatten()
702            .map(|x| x.unwrap())
703            .inspect(move |_| {
704                count.fetch_add(1, Ordering::Relaxed);
705            })
706            .collect::<Vec<String>>()
707            .await
708    }
709}
710
711impl CountReceiver<Event> {
712    pub fn receive_events<S>(stream: S) -> CountReceiver<Event>
713    where
714        S: Stream<Item = Event> + Send + 'static,
715    {
716        CountReceiver::new(|count, tripwire, connected| async move {
717            connected.send(()).unwrap();
718            stream
719                .take_until(tripwire)
720                .inspect(move |_| {
721                    count.fetch_add(1, Ordering::Relaxed);
722                })
723                .collect::<Vec<Event>>()
724                .await
725        })
726    }
727}
728
729pub async fn start_topology(
730    mut config: Config,
731    require_healthy: impl Into<Option<bool>>,
732) -> (RunningTopology, ShutdownErrorReceiver) {
733    config.healthchecks.set_require_healthy(require_healthy);
734    RunningTopology::start_init_validated(config, Default::default())
735        .await
736        .unwrap()
737}
738
739/// Collect the first `n` events from a stream while a future is spawned
740/// in the background. This is used for tests where the collect has to
741/// happen concurrent with the sending process (ie the stream is
742/// handling finalization, which is required for the future to receive
743/// an acknowledgement).
744pub async fn spawn_collect_n<F, S>(future: F, stream: S, n: usize) -> Vec<Event>
745where
746    F: Future<Output = ()> + Send + 'static,
747    S: Stream<Item = Event>,
748{
749    // TODO: Switch to using `select!` so that we can drive `future` to completion while also driving `collect_n`,
750    // such that if `future` panics, we break out and don't continue driving `collect_n`. In most cases, `future`
751    // completing successfully is what actually drives events into `stream`, so continuing to wait for all N events when
752    // the catalyst has failed is.... almost never the desired behavior.
753    let sender = tokio::spawn(future);
754    let events = collect_n(stream, n).await;
755    sender.await.expect("Failed to send data");
756    events
757}
758
759/// Collect all the ready events from a stream after spawning a future
760/// in the background and letting it run for a given interval. This is
761/// used for tests where the collect has to happen concurrent with the
762/// sending process (ie the stream is handling finalization, which is
763/// required for the future to receive an acknowledgement).
764pub async fn spawn_collect_ready<F, S>(future: F, stream: S, sleep: u64) -> Vec<Event>
765where
766    F: Future<Output = ()> + Send + 'static,
767    S: Stream<Item = Event> + Unpin,
768{
769    let sender = tokio::spawn(future);
770    tokio::time::sleep(Duration::from_secs(sleep)).await;
771    let events = collect_ready(stream).await;
772    sender.await.expect("Failed to send data");
773    events
774}
775
776#[cfg(test)]
777mod tests {
778    use std::{
779        sync::{Arc, RwLock},
780        time::Duration,
781    };
782
783    use super::retry_until;
784
785    // helper which errors the first 3x, and succeeds on the 4th
786    async fn retry_until_helper(count: Arc<RwLock<i32>>) -> Result<(), ()> {
787        if *count.read().unwrap() < 3 {
788            let mut c = count.write().unwrap();
789            *c += 1;
790            return Err(());
791        }
792        Ok(())
793    }
794
795    #[tokio::test]
796    async fn retry_until_before_timeout() {
797        let count = Arc::new(RwLock::new(0));
798        let func = || {
799            let count = Arc::clone(&count);
800            retry_until_helper(count)
801        };
802
803        retry_until(func, Duration::from_millis(10), Duration::from_secs(1)).await;
804    }
805}