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
// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.

//! Useful functions for custom test framework.
//!
//! See https://doc.rust-lang.org/unstable-book/language-features/custom-test-frameworks.html.

use crate::test::*;
use std::cell::RefCell;
use std::env;

/// A runner function for running general tests.
pub fn run_tests(cases: &[&TestDescAndFn]) {
    run_test_with_hook(cases, Nope)
}

/// Hooks for customize tests procedure.
pub trait TestHook {
    /// Called before running every case.
    fn setup(&mut self);
    /// Called after every case.
    fn teardown(&mut self);
}

/// A special TestHook that does nothing.
#[derive(Clone)]
struct Nope;

impl TestHook for Nope {
    fn setup(&mut self) {}
    fn teardown(&mut self) {}
}

struct CaseLifeWatcher<H: TestHook> {
    name: String,
    hook: H,
}

impl<H: TestHook + Send + 'static> CaseLifeWatcher<H> {
    fn new(name: String, mut hook: H) -> CaseLifeWatcher<H> {
        debug!("case start"; "name" => &name);
        hook.setup();
        CaseLifeWatcher { name, hook }
    }
}

impl<H: TestHook> Drop for CaseLifeWatcher<H> {
    fn drop(&mut self) {
        self.hook.teardown();
        debug!("case end"; "name" => &self.name);
    }
}

/// Connects std tests and custom test framework.
pub fn run_test_with_hook(cases: &[&TestDescAndFn], hook: impl TestHook + Send + Clone + 'static) {
    crate::setup_for_ci();
    let cases: Vec<_> = cases
        .iter()
        .map(|case| {
            let name = case.desc.name.as_slice().to_owned();
            let h = hook.clone();
            let f = match case.testfn {
                TestFn::StaticTestFn(f) => TestFn::DynTestFn(Box::new(move || {
                    let _watcher = CaseLifeWatcher::new(name, h);
                    f();
                })),
                TestFn::StaticBenchFn(f) => TestFn::DynTestFn(Box::new(move || {
                    let _watcher = CaseLifeWatcher::new(name, h);
                    bench::run_once(move |b| f(b));
                })),
                ref f => panic!("unexpected testfn {:?}", f),
            };
            TestDescAndFn {
                desc: case.desc.clone(),
                testfn: f,
            }
        })
        .collect();
    let args = env::args().collect::<Vec<_>>();
    test_main(&args, cases, None)
}

thread_local!(static FS: RefCell<Option<fail::FailScenario<'static>>> = RefCell::new(None));

#[derive(Clone)]
struct FailpointHook;

impl TestHook for FailpointHook {
    fn setup(&mut self) {
        FS.with(|s| {
            s.borrow_mut().take();
            *s.borrow_mut() = Some(fail::FailScenario::setup());
        })
    }

    fn teardown(&mut self) {
        FS.with(|s| {
            s.borrow_mut().take();
        })
    }
}

/// During panic, due to drop order, failpoints will not be cleared before tests exit.
/// If tests wait for a sleep failpoint, the whole tests will hang. So we need a method
/// to clear failpoints explicitly besides teardown.
pub fn clear_failpoints() {
    FS.with(|s| s.borrow_mut().take());
}

/// A runner function for running failpoint tests.
pub fn run_failpoint_tests(cases: &[&TestDescAndFn]) {
    run_test_with_hook(cases, FailpointHook)
}