mz_ore/
test.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Test utilities.
17
18use std::sync::Once;
19use std::sync::mpsc::{self, RecvTimeoutError};
20use std::thread;
21use std::time::Duration;
22
23use anyhow::bail;
24use tracing_subscriber::{EnvFilter, FmtSubscriber};
25
26static LOG_INIT: Once = Once::new();
27
28/// Initialize global logger, using the [`tracing_subscriber`] crate, with
29/// sensible defaults.
30///
31/// It is safe to call `init_logging` multiple times, and safe to call it in a
32/// binary where another test has already installed a global subscriber; in
33/// that case the installed subscriber is kept. Since `cargo test` does not run
34/// tests in any particular order, each must call `init_logging`.
35pub fn init_logging() {
36    init_logging_default("info");
37}
38
39/// Initialize global logger, using the [`tracing_subscriber`] crate.
40///
41/// The default log level will be set to the value passed in.
42///
43/// It is safe to call `init_logging_level` multiple times. Since `cargo test` does
44/// not run tests in any particular order, each must call `init_logging`.
45pub fn init_logging_default(level: &str) {
46    LOG_INIT.call_once(|| {
47        let filter = EnvFilter::try_from_env("MZ_TEST_LOG_FILTER")
48            .or_else(|_| EnvFilter::try_new(level))
49            .unwrap();
50        // Another test in this binary may have installed a global subscriber
51        // already (a turmoil test configuring its own timer, for example).
52        // `init` would panic on that, and a panic here poisons `LOG_INIT` for
53        // every `#[mz_ore::test]` in the process, so keep whatever is
54        // installed instead.
55        let _ = FmtSubscriber::builder()
56            .with_env_filter(filter)
57            .with_test_writer()
58            .try_init();
59    });
60}
61
62/// Runs a function with a timeout.
63///
64/// The provided closure is invoked on a thread. If the thread completes
65/// normally within the provided `duration`, its result is returned. If the
66/// thread panics within the provided `duration`, the panic is propagated to the
67/// thread calling `timeout`. Otherwise, a timeout error is returned.
68///
69/// Note that if the invoked function does not complete in the timeout, it is
70/// not killed; it is left to wind down normally. Therefore this function is
71/// only appropriate in tests, where the resource leak doesn't matter.
72pub fn timeout<F, T>(duration: Duration, f: F) -> Result<T, anyhow::Error>
73where
74    F: FnOnce() -> Result<T, anyhow::Error> + Send + 'static,
75    T: Send + 'static,
76{
77    // Use the drop of `tx` to indicate that the thread is finished. This
78    // ensures that `tx` is dropped even if `f` panics. No actual value is ever
79    // sent on `tx`.
80    let (tx, rx) = mpsc::channel();
81    let thread = thread::spawn(|| {
82        let _tx = tx;
83        f()
84    });
85    match rx.recv_timeout(duration) {
86        Ok(()) => unreachable!(),
87        Err(RecvTimeoutError::Disconnected) => thread.join().unwrap(),
88        Err(RecvTimeoutError::Timeout) => bail!("thread timed out"),
89    }
90}
mz_ore/test.rs

mz_ore/
test.rs
