Mocking in Async Rust

Mocking in Async Rust

There are four words in this title, and most of them deserve some kind of explanation. Let us set the scene first and look into mocking without async code.

25 October, 2023
Vortexa Analysts
Vortexa Analysts
Mocking

In this context, we’re talking about testing code, and mocking is where part of a system is replaced with “pretend” code, with two aims in mind:

  • To avoid a dependency on the real thing during testing
  • To set expectations on how we think the thing we are testing interacts with another part of the system we are not testing.

To keep our examples simple, we’re just going to define a simple trait with a couple of functions in it, and then test this in three different ways:

  1. Simple synchronous code, so we can see mocking at work
  2. Asynchronous code gone wrong, this works but it showcases how not to do it!
  3. Asynchronous testing the right, elegant way

This blog post isn’t too concerned about code composition. Just imagine though you have some code which calls something which implements the given trait, and you want to test the code you call, but mock out that trait implementation. An example might be a trait whose implementation fronts a web service. Oh and throw async into the mix!

Simple Mocking

Let’s see some code!

Ok so ExampleTrait is our example… trait. It implements two functions, one with a simple return value, one with return value wrapped in a Result.

Note this trait is decorated with the “automock” macro when building test code. This comes from the mockall crate that we’re using, and builds a mock implementation whose name is your trait’s name prefixed with Mock.

So you can see in the test code, we can actually instantiate MockExampleTrait.

Next we set some expectations. We say we expect function1() to be called once, and to return the value 10. We’re using a lambda expression so the return value could be a function of the input value, but here we chose not to bother.

Similarly we expect function2() to be called once and return Ok(11).

Now we actually call our mock object, and assert the return values are as we expect.

This example doesn’t show the “code you want to test” in the first image just to keep things as simple as possible, it is calling the mock directly. You can imagine though how you could instantiate some code of yours and pass it a mock implementation of something it uses — and then still call it in your test code.

For example, if the trait defines some “data access object” pattern for accessing a database, you could set up expectations on how you expect your code under test to interact with the database — calls it makes, and results — and then call your code without having to actually have a real database.

There’s lots of material out there on how to mock. We’ve not got to the fun bit yet.

Asynchronous mocking done badly

So I had some async code which was accessing AWS. For my unit tests, I wanted to mock this out, set expectations, and test my business logic.

The following example builds on what we did above, and tests async code. The hard way. The trait is the same, only now the two functions are async.

Here is the corresponding test code:


                        
                    

… and here is the horror boilerplate code to make that testing code work:


                        
                    

I’m sorry, what just happened?

We’re using the async-trait crate to allow us to create, well, async traits. These are not yet finalized in stable Rust, but should be coming soon. By declaring a trait with this macro, it actually turns this:

…into this:

There’s a lot going on here.

Under the hood, async functions don’t really return results, they return Futures which in turn can give us results. Futures are polled to see if they are ready to return a value. They might be blocked by an IO operation or even some computation. Different threads may be involved during this polling, so the data needs to be pinned to a fixed address on the heap. If the data type cannot be pinned in such a way, it needs to be boxed first so it can be put on the heap. It must also implement the Send trait as it can be passed between threads.

I really recommend looking at the Asynchronous Programming in Rust book for more details, of which there are plenty.

Note we’re using the #[tokio:test] macro from to support the running of asynchronous tests.

What about mocking? Well, the disaster here (which took me a while to discover, hence this blog post) is the order of the two macro calls, using async-trait before automock. Essentially, that means:

  1. Convert the code to use Futures and complex types under the hood
  2. Then add an auto-generated mock implementation based on that

You see the problem?

In the test code, when we say we expect function1() to return 10, we actually have to write that it returns pinned Future that will return 10. Future is a trait, so what do we do for an implementation? We create one. See the ConstantFuture the implementation. Even though the data is actually just a constant, we need to implement a poll() function for all the machinery to work. In our case, poll() will always return that we’re immediately ready, there is no underlying resource which can block us.

Admittedly the implementation is generic and as such clones the value, we could have made it usize specific and slightly simpler. A generic implementation is more useful however when dealing with different data types.

What about function2()? This returns a Result which has two generic parameters, the data type and the error type. See the ConstantResultFuture implementation. We actually need to use PhantomData so we don’t lose track of the error type, and can implement the correct poll() signature.

This code actually works, teaches us much about async under the hood, but is a bit… horrifying.

Doing it right

The problem above is we turned the trait into async code, and then implemented a mock object from it.

How about doing this the other way around?

  1. Auto-generate a mock implementation
  2. Then add the async support

Call the macros in the opposite order. Here’s the code:

That’s it! Simple. The lambda expressions used in the expect() calls when setting up the mock are nice and synchronous. If you want to return 10, just return 10. Or Ok(11) if you prefer. The trait still supports asynchronous calls in your production code, and the mock implementation is also asynchronous.

We can have our cake and eat it. This version uses no contrived Future implementations.

The order the macros are declared, for mocking and async traits, is vitally important.

Conclusion

Mocking is easy in Rust, but mocking async code in Rust is also very achievable given the right tools. Async code can be easily replaced and mocked-out in test code, so if your business logic must be mixed with code you need to mock away, you can. If you do it the hard way, you can learn a lot in the process too.

Vortexa Analysts
Vortexa
Vortexa Analysts