如何在测试中为无状态Trait引入状态？

问题场景 #

我开发了一个小型库，核心是用于模拟特定功能的插件，核心代码如下：

trait ProvideFunctions {
    fn do_a();
    fn do_b(foo: usize) -> ();
}

struct SomeThing<P: ProvideFunctions> {
    _marker: PhantomData<P>
}

struct MainThing<P: ProvideFunctions> {
    // some fields
    some_thing: SomeThing<P>,
    _marker: PhantomData<P>
}

impl<P: ProvideFunctions> MainThing<P> {
   fn do_something(&self) -> () {
        P::do_a();
        P::do_b(1);
        // ...
   }
}

这些插件函数设计为无状态，但编写测试时存在不便——测试中状态十分有用（如跟踪调用、构造模拟对象等）。

我想到两种方案，但都存在问题：

• 为测试使用static状态，但这并不合适，因为状态仅属于特定测试，且测试并行运行时会出现污染。
• 将插件改为有状态，但所有共享该功能的结构体都需要持有插件副本。尽管编译器会在实际场景中优化掉状态，但Trait本身会被引导向有状态设计，偏离原有设计初衷。

请问这类问题是否有通用解决方案？还有其他可行思路吗？

通用解决方案：关联类型 + 双环境实现 #

通过给ProvideFunctions trait添加关联类型，可以在不破坏生产环境无状态设计的前提下，为测试环境提供状态支持：

调整后的核心代码 #

use std::marker::PhantomData;

trait ProvideFunctions {
    // 定义关联类型，生产环境用空类型，测试用带状态的类型
    type State: Default;
    // 方法新增状态参数，生产实现忽略该参数，测试实现利用它
    fn do_a(state: &mut Self::State);
    fn do_b(state: &mut Self::State, foo: usize);
}

// 生产环境的无状态实现：用空类型作为State
struct ProductionImpl;
impl ProvideFunctions for ProductionImpl {
    type State = ();
    fn do_a(_state: &mut ()) {
        // 原有的无状态业务逻辑
    }
    fn do_b(_state: &mut (), foo: usize) {
        // 原有的无状态业务逻辑
    }
}

// 调整结构体，关联State类型
struct SomeThing<P: ProvideFunctions> {
    _marker: PhantomData<P>
}

struct MainThing<P: ProvideFunctions> {
    some_thing: SomeThing<P>,
    // 生产环境中State是()，编译器会自动优化掉这个字段
    state: P::State,
    _marker: PhantomData<P>
}

impl<P: ProvideFunctions> MainThing<P> {
    // 构造函数：生产环境自动生成空状态，测试环境可自定义状态
    fn new() -> Self {
        Self {
            some_thing: SomeThing { _marker: PhantomData },
            state: P::State::default(),
            _marker: PhantomData
        }
    }

    fn do_something(&mut self) {
        P::do_a(&mut self.state);
        P::do_b(&mut self.state, 1);
    }
}

测试环境的带状态实现 #

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

    // 测试用状态结构体，跟踪调用次数和参数
    struct TestState {
        do_a_calls: usize,
        do_b_args: Vec<usize>,
    }

    impl Default for TestState {
        fn default() -> Self {
            Self {
                do_a_calls: 0,
                do_b_args: Vec::new(),
            }
        }
    }

    struct TestImpl;
    impl ProvideFunctions for TestImpl {
        type State = TestState;
        fn do_a(state: &mut TestState) {
            state.do_a_calls += 1;
        }
        fn do_b(state: &mut TestState, foo: usize) {
            state.do_b_args.push(foo);
        }
    }

    #[test]
    fn test_do_something() {
        let mut main_thing = MainThing::<TestImpl>::new();
        main_thing.do_something();
        
        assert_eq!(main_thing.state.do_a_calls, 1);
        assert_eq!(main_thing.state.do_b_args, vec![1]);
    }
}

方案优势 #

• 生产环境完全保留无状态本质：State类型为()，编译器会彻底优化掉state字段和相关参数，性能与原代码一致
• 测试环境可灵活实现状态跟踪，每个测试实例拥有独立状态，避免测试污染
• 仅对Trait做最小扩展，无需大幅改动原有结构体设计，保持API兼容性

轻量替代方案：线程局部存储（Thread-Local） #

如果不想修改原有Trait定义，可以用线程局部存储为测试提供隔离状态：

#[cfg(test)]
thread_local! {
    static TEST_STATE: std::cell::RefCell<TestState> = std::cell::RefCell::new(TestState::default());
}

#[cfg(test)]
struct TestImpl;
#[cfg(test)]
impl ProvideFunctions for TestImpl {
    fn do_a() {
        TEST_STATE.with(|s| s.borrow_mut().do_a_calls += 1);
    }
    fn do_b(foo: usize) {
        TEST_STATE.with(|s| s.borrow_mut().do_b_args.push(foo));
    }
}

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

    struct TestState {
        do_a_calls: usize,
        do_b_args: Vec<usize>,
    }

    impl Default for TestState {
        fn default() -> Self {
            Self {
                do_a_calls: 0,
                do_b_args: Vec::new(),
            }
        }
    }

    #[test]
    fn test_do_something() {
        // 测试前重置状态，避免跨测试污染
        TEST_STATE.with(|s| *s.borrow_mut() = TestState::default());
        
        let main_thing = MainThing::<TestImpl>::new();
        main_thing.do_something();
        
        TEST_STATE.with(|s| {
            let state = s.borrow();
            assert_eq!(state.do_a_calls, 1);
            assert_eq!(state.do_b_args, vec![1]);
        });
    }
}

注意事项 #

• 线程局部存储保证了不同测试线程的状态隔离，但单线程运行时需手动重置状态
• 无需修改原有Trait和结构体，侵入性极小
• 生产环境完全不受影响，测试代码仅在#[cfg(test)]下编译

内容的提问来源于stack exchange，提问作者milck