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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
//! Miscellaneous test code for Zebra.
#![doc(html_favicon_url = "https://zfnd.org/wp-content/uploads/2022/03/zebra-favicon-128.png")]
#![doc(html_logo_url = "https://zfnd.org/wp-content/uploads/2022/03/zebra-icon.png")]
#![doc(html_root_url = "https://docs.rs/zebra_test")]
// Each lazy_static variable uses additional recursion
#![recursion_limit = "512"]

use std::sync::Once;

use color_eyre::section::PanicMessage;
use once_cell::sync::Lazy;
use owo_colors::OwoColorize;
use tracing_error::ErrorLayer;
use tracing_subscriber::{fmt, prelude::*, EnvFilter};

#[allow(missing_docs)]
pub mod command;

pub mod mock_service;
pub mod net;
pub mod network_addr;
pub mod prelude;
pub mod service_extensions;
pub mod transcript;
pub mod vectors;
pub mod zip0143;
pub mod zip0243;
pub mod zip0244;

/// A single-threaded Tokio runtime that can be shared between tests.
/// This runtime should be used for tests that need a single thread for consistent timings.
///
/// This shared runtime should be used in tests that use shared background tasks. An example is
/// with shared global `Lazy<BatchVerifier>` types, because they spawn a background task when they
/// are first initialized. This background task is stopped when the runtime is shut down, so having
/// a runtime per test means that only the first test actually manages to successfully use the
/// background task. Using the shared runtime allows the background task to keep running for the
/// other tests that also use it.
///
/// A shared runtime should not be used in tests that need to pause and resume the Tokio timer.
/// This is because multiple tests might be sharing the runtime at the same time, so there could be
/// conflicts with pausing and resuming the timer at incorrect points. Even if only one test runs
/// at a time, there's a risk of a test finishing while the timer is paused (due to a test failure,
/// for example) and that means that the next test will already start with an incorrect timer
/// state.
pub static SINGLE_THREADED_RUNTIME: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("Failed to create Tokio runtime")
});

/// A multi-threaded Tokio runtime that can be shared between tests.
/// This runtime should be used for tests that spawn blocking threads.
///
/// See [`SINGLE_THREADED_RUNTIME`] for details.
pub static MULTI_THREADED_RUNTIME: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("Failed to create Tokio runtime")
});

static INIT: Once = Once::new();

/// Initialize global and thread-specific settings for tests,
/// such as tracing configs, panic hooks, and `cargo insta` settings.
///
/// This function should be called at the start of every test.
///
/// It returns a drop guard that must be stored in a variable, so that it
/// gets dropped when the test finishes.
#[must_use]
pub fn init() -> impl Drop {
    // Per-test

    // Settings for threads that snapshots data using `insta`

    let mut settings = insta::Settings::clone_current();
    settings.set_prepend_module_to_snapshot(false);
    let drop_guard = settings.bind_to_scope();

    // Globals

    INIT.call_once(|| {
        let fmt_layer = fmt::layer().with_target(false);
        // Use the RUST_LOG env var, or by default:
        //  - warn for most tests, and
        //  - for some modules, hide expected warn logs
        let filter_layer = EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| {
                // These filters apply when RUST_LOG isn't set
                EnvFilter::try_new("warn")
                    .unwrap()
                    .add_directive("zebra_consensus=error".parse().unwrap())
                    .add_directive("zebra_network=error".parse().unwrap())
                    .add_directive("zebra_state=error".parse().unwrap())
                    .add_directive("zebrad=error".parse().unwrap())
                    .add_directive("tor_circmgr=error".parse().unwrap())
            })
            // These filters apply on top of RUST_LOG.
            // Avoid adding filters to this list, because users can't override them.
            //
            // (There are currently no always-on directives.)
            ;

        tracing_subscriber::registry()
            .with(filter_layer)
            .with(fmt_layer)
            .with(ErrorLayer::default())
            .init();

        color_eyre::config::HookBuilder::default()
            .add_frame_filter(Box::new(|frames| {
                let mut displayed = std::collections::HashSet::new();
                let filters = &[
                    "tokio::",
                    "<futures_util::",
                    "std::panic",
                    "test::run_test_in_process",
                    "core::ops::function::FnOnce::call_once",
                    "std::thread::local",
                    "<core::future::",
                    "<alloc::boxed::Box",
                    "<std::panic::AssertUnwindSafe",
                    "core::result::Result",
                    "<T as futures_util",
                    "<tracing_futures::Instrumented",
                    "test::assert_test_result",
                    "spandoc::",
                ];

                frames.retain(|frame| {
                    let loc = (frame.lineno, &frame.filename);
                    let inserted = displayed.insert(loc);

                    if !inserted {
                        return false;
                    }

                    !filters.iter().any(|f| {
                        let name = if let Some(name) = frame.name.as_ref() {
                            name.as_str()
                        } else {
                            return true;
                        };

                        name.starts_with(f)
                    })
                });
            }))
            .panic_message(SkipTestReturnedErrPanicMessages)
            .install()
            .unwrap();
    });

    drop_guard
}

/// Initialize globals for tests that need a separate Tokio runtime instance.
///
/// This is generally used in proptests, which don't support the `#[tokio::test]` attribute.
///
/// If a runtime needs to be shared between tests, use the [`SINGLE_THREADED_RUNTIME`] or
/// [`MULTI_THREADED_RUNTIME`] instances instead.
///
/// See also the [`init`] function, which is called by this function.
pub fn init_async() -> (tokio::runtime::Runtime, impl Drop) {
    let drop_guard = init();

    (
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("Failed to create Tokio runtime"),
        drop_guard,
    )
}

struct SkipTestReturnedErrPanicMessages;

impl PanicMessage for SkipTestReturnedErrPanicMessages {
    fn display(
        &self,
        pi: &std::panic::PanicHookInfo<'_>,
        f: &mut std::fmt::Formatter<'_>,
    ) -> std::fmt::Result {
        let payload = pi
            .payload()
            .downcast_ref::<String>()
            .map(String::as_str)
            .or_else(|| pi.payload().downcast_ref::<&str>().cloned())
            .unwrap_or("<non string panic payload>");

        // skip panic output that is clearly from tests that returned an `Err`
        // and assume that the test handler has already printed the value inside
        // of the `Err`.
        if payload.contains("the test returned a termination value with a non-zero status code") {
            return write!(f, "---- end of test output ----");
        }

        writeln!(f, "{}", "\nThe application panicked (crashed).".red())?;

        write!(f, "Message:  ")?;
        writeln!(f, "{}", payload.cyan())?;

        // If known, print panic location.
        write!(f, "Location: ")?;
        if let Some(loc) = pi.location() {
            write!(f, "{}", loc.file().purple())?;
            write!(f, ":")?;
            write!(f, "{}", loc.line().purple())?;
        } else {
            write!(f, "<unknown>")?;
        }

        Ok(())
    }
}