🍐

Rust グローバル変数をテストで隔離する方法

2023/05/09に公開1

通常、開発中でグローバルな状態を作成するために、lazy_static または once_cell などのライブラリが使用されます。以下は例です。

use lazy_static::lazy_static;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

pub struct Registry {
    total: AtomicUsize,
}

impl Registry {
    pub fn new() -> Self {
        Self {
            total: AtomicUsize::default(),
        }
    }

    pub fn register(&self, label: &str) {
        self.total.fetch_add(1, Ordering::SeqCst);
    }

    pub fn get_total(&self) -> usize {
        self.total.load(Ordering::SeqCst)
    }
}

lazy_static! {
    static ref DEFAULT_REGISTRY: Registry = {
        let r = Registry::new();

        r
    };
}

pub fn register_xxx() {
    // do something
    DEFAULT_REGISTRY.register("xxx");
}

pub fn register_yyy() {
    // do something
    DEFAULT_REGISTRY.register("yyy");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_1() {
        let registry = &DEFAULT_REGISTRY;
        registry.register("sample");
        assert_eq!(registry.get_total(), 1);
    }

    #[test]
    fn test_2() {
        let registry = &DEFAULT_REGISTRY;
        assert_eq!(registry.get_total(), 0);
    }

    #[test]
    fn test_send() {
        fn test<C: Send>() {}
        test::<DEFAULT_REGISTRY>();
    }

    #[test]
    fn test_sync() {
        fn test<C: Sync>() {}
        test::<DEFAULT_REGISTRY>();
    }
}

この例では、handler を登録するためのグローバルな DEFAULT_REGISTRY を定義しました。しかし、単体テストを書く際に問題が発生する可能性があります。Rust の単体テストは、各テストケースを並行して実行するため、同じ DEFAULT_REGISTRY にアクセスすることになり、テストケース間で書き込まれたデータが相互に影響し合います。以下は、単体テストの実行後の結果です。

ここにいくつかの解決策があります。

方法 1・ テスト中にグローバル変数にアクセスしない

例として、テストケースを書くときに、default_registry ではなく手動で Registry::new() を使用する。見た目は奇妙かもしれないが、根本的な問題を解決することができます。

    #[test]
    fn test_1() {
        let registry = Registry::new();
        registry.register("sample");
        assert_eq!(registry.get_total(), 1);
    }

この方法の制限は、現在のコードが default_registry と深く結合している場合、使用できないことです。例えば、上記のコードでは register_xxx 関数が DEFAULT_REGISTRY を参照しています。そのため、テストケースでこの振る舞いを変更することができません。

方法 2・条件付きコンパイル

DEFAULT_REGISTRY へのアクセスを制御し、条件付きコンパイルを使用して通常の実行とユニットテストの実行フローを変更します。

lazy_static! {
    static ref DEFAULT_REGISTRY: Registry = {
        let r = Registry::new();

        r
    };
}

#[cfg(not(test))]
pub fn default_registry() -> &'static Registry {
    lazy_static::initialize(&DEFAULT_REGISTRY);
    &DEFAULT_REGISTRY
}

#[cfg(test)]
pub fn default_registry() -> &'static Registry {
    let r = Registry::new();
    let b = Box::new(r);
    let static_ref: &'static mut Registry = Box::leak(b);
    static_ref
}

すべての DEFAULT_REGISTRY への直接的アクセス箇所を default_registry() に変更し、cfg を使用して異なる関数をコンパイルすることにより、DEFAULT_REGISTRY のアクセスを制御します。ユニットテストの場合、新しい registry オブジェクトを作成し、Box::leak を使用してライフタイムを static に昇格させ、非テスト条件下の関数シグネチャと一致するようにします。ただし、この方法にも避けられない問題があります。Box::leak は、実行時の値を static に変換します。これは、実際には Box のスタック上での解放に責任を負っていた部分を忘れ、ヒープ上に割り当てられたBox の部分のポインタを返すことと同じです。leak と forget は、実際には ManuallyDrop を使用して実装されています。したがって、このメモリブロックを手動で解放する必要がありますが、これは単体テストの一部であるため、まずはこの問題を考慮しなくても大丈夫です。

同じスレッド内で何度も default_registry を呼び出した場合、異なる Registry オブジェクトが取得されます。そのため、以下のようにコーディングすることが制約されることになります。

let registry = default_registry();

// NO
default_registry().register("sample");
// OK
registry.register("sample");  

assert_eq!(registry.get_total(), 1);

方法 3・Thread Local (TLS)

前述の2つの方法を組み合わせて、Thread local を利用して新しいバージョンを実装することができます。

cfg_if!(
if #[cfg(not(test))]
{
    lazy_static! {
        static ref DEFAULT_REGISTRY: Registry = Registry::new();
    }
}
else
{
    use std::cell::Cell;

    thread_local!{
        static LOCAL_REGISTRY: Cell<Option<&'static Registry>> = Cell::new(None);
    }

    struct RegistryProxy;

    impl std::ops::Deref for RegistryProxy {
        type Target = Registry;

        #[inline]
        fn deref (&self) -> &Self::Target {
            LOCAL_REGISTRY.with(|cl| {
                if cl.get().is_none() {
                    let r = Registry::new();
                    let b = Box::new(r);
                    let static_ref = Box::leak(b);
                    cl.set(Some(static_ref));
                }
                cl.get().unwrap()
            }
            )
        }
    }

    static DEFAULT_REGISTRY: RegistryProxy = RegistryProxy;
});

まず、第2の方法に従って、 Deref trait を使用して dereferencing の動作を変更します。まず、 Cell を使用して内部可変性を与え、thread local 変数(TLS)を作成します。初期値は None で、誰もここから Registry を取得していないことを示します。次に、 RegistryProxy を実装し、 deref でこの TLS が初期化されたかどうかを確認し、初期化されていない場合は新しいオブジェクトを作成し、初期化されている場合は既存のオブジェクトを返します。

ここでの DEFAULT_REGISTRY は実際には RegistryProxy 型であり、Registry ではありませんが、自動的なデリファレンスにより、動作は Registry と同じです。

方法 4・DI、Mock

Java や Python などの言語で一般的ですが、Rust ではやや複雑です。mockall などのフレームワークを使用して、実装することができますが、コードの変更量は多いです。

Discussion

Hanaasagi's JP BlogHanaasagi's JP Blog

std::sync::Once version

    use std::cell::Cell;
    use std::sync::Once;

    thread_local!{
        static LOCAL_REGISTRY: (Cell<Option<&'static Registry>>, Once) = (Cell::new(None), Once::new());
    }

    struct RegistryProxy;
    impl std::ops::Deref for RegistryProxy {
        type Target = Registry;

        #[inline]
        fn deref(&self) -> &Self::Target {
            LOCAL_REGISTRY.with(|cl| {
                cl.1.call_once(|| {
                    let static_ref = Box::leak(Box::new(Registry::new()));
                    cl.0.set(Some(static_ref));
                });

                unsafe {
                    match *cl.0.as_ptr() {
                        Some(ref x) => x,
                        None => {
                            unreachable!()
                        }
                    }
                }
            })
        }
    }