All the ways to mock your Rust code

发表于 更新于 温馨的阅读时光

本文来源

起因

你试图用rust写一个k8s控制器，你意识到直接写完出错了不好诊断，你决定emit一些事件以获取一些上下文，你运行了所有测试并期待通过，你认为这是一定的因为你什么都没改……

但是测试挂了，因为单元测试里根本就没有运行着k8s server，请求根本收不到

What do???

以上正是本文作者在做Simkube项目时遇到的实际问题，作者花了大量的时间探索一个测试的最佳实践，于是有了这篇博客：all the ways to mock a network call in Rust

本文也触及了作者在写代码时最烦的一点:

“DON’T MAKE ME MAKE MY PRODUCTION CODE WORSE JUST TO MAKE IT EASIER TO TEST, COMMA JERK”
（“千万别逼我为了方便测试，就把老子的生产代码改得一团糟，【此处顿号】你个混蛋！” Gemini 译）

背景

kube-rs库提供了内置的写controller的接口，用户需要提供的主要是一个reconcile函数，它将在无论何时controller试图做点什么时被调用，签名像这样：

pub async fn reconcile<T: kube::Resource>(
    object: Arc<T>, global_ctx: Arc<GlobalContext>) -> Result<Action, Error>

第一个参数是k8s对象，注明预期的k8s状态是什么，第二个参数是一个上下文用于获取当前k8s状态，本文里大概是这样：

struct GlobalContext {
    client: kube::Client,
    recorder: SkEventRecorder,
}

一个连接k8s服务端的客户端和一个事件记录器，对于recorder，一个简化的实现如下：

use kube::runtime::events;

struct SkEventRecorder {
    recorder: events::Recorder,
}

impl SkEventRecorder {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }

    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
}

一个kube::runtime::events::Recorder和一些辅助方法。这里最终涉及了网络请求，每一次publish都把事件发给api server并期待一个回应。

那么如何把这段代码写成可测试的样子呢？

过程

Method 1: Boxed Trait

最常用的方法，在我的上一篇博客里亦有记载，使用接口和依赖注入，在生产环境使用一个真实的recorder，在测试时使用一个mock的

于是写起来类似于这样：

#[async_trait]
pub trait EventRecordable {
    pub async fn send_event_A(&self) -> Result<(), Error>;
    pub async fn send_event_B(&self) -> Result<(), Error>;
}

struct SkEventRecorder {
    recorder: events::Recorder,
}

impl SkEventRecorder {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }
}

impl EventRecordable for SkEventRecorder {
    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
}

struct GlobalContext {
    client: kube::Client,
    recorder: Arc<dyn EventRecordable>,
}

这个写法类似于go语言，在测试时只需要定义一个MockSkEventRecorder，实现一个EventRecordable特征就可以完美的嵌入，实际使用的recorder的类型不再是一个具体类型，而是dyn EventRecordable

由于实际上需要mock的只有recorder，send_event方法可以复用，所以也可以写成这样：

#[async_trait]
trait EventRecordable {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error>;
}

struct EventRecorder {
    recorder: events::Recorder,
}

impl EventRecordable for EventRecorder {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error> {
        Ok(self.recorder.publish(event_a).await?)
    }
}

struct SkEventRecorder {
  recorder: Arc<impl EventRecordable>,
}

impl SkEventRecorder {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }

    #[cfg(test)]
    pub fn new_with_mock() -> SkEventRecorder { 
        /* initialize a new recorder with a mock sender */ 
    }

    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
}

struct GlobalContext {
    client: kube::Client,
    recorder: SkEventRecorder,
}

使用trait的位置往下放了一层，优点是降低了代码量，但也有缺点：

  • send_event方法是async的，所以EventRecordable 特征是async trait，但这不是默认就支持的，需要额外引入
  • 引入了额外的抽象只是为了易于测试，为阅读代码的人增加了不必要的心智负担
  • 影响性能！！！原本类型是编译器可知的，现在变成了一个指针，编译器不知道impl EventRecordable这个结构体有多大，于是只能把它分配在堆上。因为代码是高度异步的，所以还有一层Arc，这就造成了一次指针查找和一次Rc操作的开销。同时为了找到定义的方法，还有一步动态分发（vtable 或其他）

也许写go的觉得这些不需要关心，但是，作者说，这是rust！这不能接受

Generics

动态分发可以通过泛型参数变成静态的，看起来像这样：

#[async_trait]
pub trait EventRecordable {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error>;
}

struct EventRecorder {
    recorder: events::Recorder,
}

impl EventRecordable for EventRecorder {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error> {
        Ok(self.recorder.publish(event_a).await?)
    }
}

struct SkEventRecorder<T: EventRecordable> {
    recorder: T,
}

impl<T: EventRecordable> SkEventRecorder<T> {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }

    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
}

struct GlobalContext<T: EventRecordable> {
    client: kube::Client,
    recorder: SkEventRecorder<T>,
}

这份代码把运行时开销转移到了编译时，于是不会有堆开销，不会有vtable查找开销

但是依然有缺点—— $#&% 这样的东西太丑了，这个愚蠢的模板参数会一路传播到顶层的GlobalContext里，其他人来读这份代码又需要额外理解很多东西，这也是不可以接受的！！！

Compile-time feature flags (take 1)

如果……不用这个trait实现呢？

struct SkEventRecorder {
    recorder: events::Recorder,
}

impl SkEventRecorder {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }
}

#[cfg(not(test))]
impl SkEventRecorder {
    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.publish(event_a).await?)
    }
}

#[cfg(test)]
impl SkEventRecorder {
    pub async fn send_event_A(&self) -> Result<(), Error> { Ok(()) }
    pub async fn send_event_B(&self) -> Result<(), Error> { Ok(()) }
}

struct GlobalContext {
    client: kube::Client,
    recorder: SkEventRecorder,
}

利用编译时flag就可以把代码写的cute一点。在正式运行时，走一条编译分支，在测试时又走另一边。

但是有两个问题：

  • 如果想要测试复杂的逻辑，这个模式不便于拓展
  • 在实际场景中，事件记录器和上下文可能定义在不同的crate里，但是test的生效范围以crate为界，为了能把test的flag传过去又需要很多工作，糟透了

Method 4: Compile-time feature flags (take 2)

把两种模式结合一下呢：

#[async_trait]
pub trait EventRecordable {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error>;
}

struct EventRecorder {
    recorder: events::Recorder,
}

impl EventRecordable for EventRecorder {
    async fn send_event(&self, event: &events::Event) -> Result<(), Error> {
        Ok(self.recorder.publish(event_a).await?)
    }
}

struct SkEventRecorder {
    #[cfg(not(test))]
    recorder: EventRecorder,

    #[cfg(test)]
    recorder: MockEventRecorder
}

impl SkEventRecorder {
    pub fn new() -> SkEventRecorder { /* initialize a new recorder */ }

    pub async fn send_event_A(&self) -> Result<(), Error> {
        let event_a = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
    pub async fn send_event_B(&self) -> Result<(), Error> {
        let event_b = events::Event::new( /* make an event */ );
        Ok(self.recorder.send_event(event_a).await?)
    }
}

struct GlobalContext {
    client: kube::Client,
    recorder: SkEventRecorder,
}

这个frankentrait有cfg(test)的所有缺点，看起来也愚蠢

Method 5: The actor framework

既然以上方法都很糟，那就先回顾一下这里的代码究竟做了什么事：
发送一些事件 通过网络 然后测试

一个也许是标准的实现是用Actor Model ，基本思想是，自己的程序代码是actor，彼此之间只通过发送message交流（通过某种队列）

如果用这个模型，那么就需要重写reconcile 函数然后只把事件发送到一个队列里，然后写一些代码在另一端去异步监听这个队列并把事件发送到K8s apiserver通过网络

Actor模型显然很好也有很多应用，但是这段代码只是想记录一下事件，为此影响项目架构是不合适的

Method 6: Fake the apiserver

核心的问题在于，以上所有方法都需要把生产代码从某种简单直接的东西变成**只为了易于测试，那么也许可以换个视角……保持代码不变，然后mock 一个 K8s 的apiserver呢？这样就可以用真实的api server做真实的交互，只有一个小问题，需要写一个K8s apiserver，然后就可以通过所有的单元测试了。以及我们在单元测试中引入了网络调用，众所周知，网络可是出了名地极其健壮，从不丢包，绝对不会胡乱造成不稳定的测试（flaky tests），绝——对——不——会。

Method 7: There is no method 7

没有然后了，所有方法都不能让人满意。强类型编译型语言需要你告诉编译器你有一个具体的类型做一件具体的事，在编译期，不幸的是这就是会使测试变得困难

那么作者最后选择了哪种呢？

claude:“What are you doing, you dumbass? The cost of a network request is orders of magnitude (几个数量级) higher than the cost of a couple of pointer lookups. Just stick it in an Arc and call it a day.”

作者：“So that’s what I did. I’m kindof lowkey upset that the stochastic labubu (随机labubu，指AI) was right about this, but, well, I guess we can’t win them all.”

# rust# test