Testing with Typestates

Most tests I've read or written are a big blob of code with a few assertions wedged in. They're what we're taught in school and the style is ubiquitous in Java and other OOP languages and are considered the "default" way to write a test. But I wanted to try something different for my API test suite. Why?

There are a few problems with these tests:

Here's our example test:

#[test]
get_post_by_title() {
let test_app = create_test_app().await;
let client = test_app.client;

let post: PostData = Faker.fake();

let submit_response = submit_post_helper(post.clone(), client).await;
assert_eq!(StatusCode::OK, submit_response.status);

let fetch_response = get_post_by_title_helper(post.title.clone()).await;
assert_eq!(StatusCode::OK, fetch_response.status);

assert_eq!(post.title.clone(), fetch_response.json.title())

let bad_fetch_response = get_post_by_title_helper(post.title.clone()).await;
assert_eq!(StatusCode::NOT_FOUND, bad_fetch_response.status);
}

At this point there are 6 variables in scope and we're writing and rewriting a lot of similar code despite our helpers. If this is as complicated as your tests get, this is just fine! But in my experience it snowballs from here, getting more error-prone, verbose, and impenetrable.

What I want is something that feels like the Rust Iterator experience but for tests. A series of modular, scoped operations which are abstracted out from the problem space so that you can focus on business logic. The basic flow might look like this:

I'd heard that the typestate pattern could be good for writing tests so I decided to start there. It worked better than I could have hoped! Here's the version of the same test that we'll work towards today.

#[tokio::test]
async fn get_post_by_title() {
let post: PostData = Faker.fake();
let bad_title = PostTitle::new("Title that doesn't exist in the DB");

Test::new()
.take(post.clone())
.submit_post::<Ok>()
.await
.map(|state| state.title)
.get_post_from_title::<Ok>()
.await
.inspect(|state| assert_eq!(state, &post, "Fetched the wrong post"))
.take(bad_title)
.get_post_from_title::<ClientError>()
.await
.cleanup();
}

We've separated the test variables from the logic and we can now read line by line a simple list of steps the test is performing. Each member method of the Test struct can be reused in any other test with different data.

Our Test struct includes an instance of our app and one additional field, state. The type of state dictates the methods that can be called on our Test. What's really nice about this when writing tests is that as soon as you hit that . for the next test step, your editor will show you a list of valid actions to take next.

Much of this is a textbook example of the typestate pattern - with the exception of those generics: Ok & ClientError. What's up with those?

Well, we want to assert the expected behavior of the REST request - ideally without duplicating our test methods. And a get_post call should permit the test to move on to a following stage, even if the response was bad, as long as that's what was anticipated.

We want to avoid a Test<Option<PostData>> state - each stage should be able to act on the state with absolute certainty that the state for that stage exists. Otherwise we have to constantly check our Options and we'll be prone to causing test failures in places other than where they originated. I'm looking at you, Java.

impl Test<PostData> {
fn replace_data_body(mut self, new_body: String) -> Test<PostData> {
self.data.body = new_body;
self
}
}

We don't have to check if self.data is present. We don't have to account for all the outcomes of the previous test stage. This method can't be called on a version of Test where data is not of the type PostData.

Here's how we can use associated types to return the Output type on the impl Expectation whenever we call a function that makes a REST request.

trait Expectation<T> {
type Output;
fn resolve(response: Response);
}

struct Ok;
impl<T> Expectation<T> for Ok {
type Output = T;
fn resolve(response: Response) -> Self::Output {
// ...
}
}

struct ClientError;
impl<T> Expectation<T> for ClientError {
// rest calls with this expectation will yield `Test<()>` to the next stage.
type Output = ();
}

impl Test {
fn get_request<Exp: Expectation>(self, post_id: String) -> Test<Exp::Output> {
let response = self.client.get_request(...).await.unwrap();

Exp::resolve(response) // expectation-dependent resolution
}
}

When default associated types become stable we'll be able to get rid of some of the boilerplate of setting up different Expectation types.

Some utility functions I wrote include take, which imports new state into the Test instance, map, for mapping the test state into another type, and fork, for cloning the whole Test instance and following a branching execution path inside a closure. This can help out with situations where you want to test some child fields of your state before returning to the full state for the next test steps. Finally, inspect is good for injecting assertions into your test without writing a custom test stage for just that purpose.

Let's look again at our final test code, with the type annotations you'll see in an editor:

#[tokio::test]
async fn get_post_by_title() {
let post: PostData = Faker.fake();

Test::new()
.take(post.clone()) // Test<PostData>
.submit_post::<Ok>()
.await // Test<PostData>
.map(|state| state.title) // Test<PostTitle>
.get_post_from_title::<Ok>()
.await // Test<PostData>
.assert_eq(post.clone(), "Fetched the wrong post")
.take(PostTitle::new("Title that doesn't exist in the DB")) // Test<PostTitle>
.get_post_from_title::<ClientError>()
.await // Test<()> - expected an error so no post returned
.cleanup();
}

There's more work to be done here but I'm liking how it's turning out so far. For example, the Drop trait should be implemented for the test context so that any databases get dropped should the test fail due to a panic. I'll revisit this later.