Implementing a network protocol with Swift and snapshot testing

A guy standing next to a beanstalk

In a previous post, I mentioned Beanstalk, a simple and fast work queue. It uses a simple and straightforward protocol, and learning to implement it in Swift has proven to be extremely fun.

Why Beanstalk?

A couple of years ago, I started dabbling with RabbitMQ, a message broker that my team and I used to process large amounts of data in the background. This experience clearly left a mark, which is probably why I have this strange fascination for work queues now.

Recently, I came across Beanstalkd and read the protocol implementation. It communicates using ASCII text over TCP, so the exchanges between the server and the client are very easy to comprehend. I figured I could set myself a personal challenge and implement this protocol in Swift. Without further ado, let’s see how to queue a new job into a beanstalk server.

Describing a message

Contrary to other protocols like AMQP that is binary, the beanstalk protocol uses ASCII, which makes it very easy to understand: the messages sent to the server are human-readable. For example, one of the messages you’ll end up sending the most is the put command, which queues a new job.

For each put command, the beanstalk server expects the following information:

  • A priority, which I covered in a previous blog post about Tagged. It’s an Int32. Jobs with a lower value will be considered more urgent than jobs with a larger value, which means that a job with a priority of “0” is the most urgent job.
  • A delay, which is a number of seconds, before marking a job as “ready”. From the moment it was queued, until that number of seconds has passed, the job will not be ready to be consumed and processed.
  • A time to run, which is the maximum number of seconds allowed for a worker to finish processing that job. After that delay, the job will time out.
  • A job “body” and the number of bytes for this body.

Once properly formatted, the payload that will be sent to the beanstalk server looks like this:

put <pri> <delay> <ttr> <bytes>\r\n<data>\r\n

In Swift, we can represent this message as the following struct:

struct PutCommand {
    let priority: Priority

    let delay: Seconds

    let timeToRun: Seconds

    let body: String
}

Note: the Priority and the Seconds types have been covered in the blog post about Tagged. If you haven’t read it yet or if you’re confused, those are basically integers, covered in type safety.

The size of the body is not present in this struct, even though it’s required by the beanstalk protocol. Since it’s dependent on the body itself, there is no point in storing it. We can calculate it when we turn this struct into a payload that the beanstalk server will understand.

In fact, let’s do this right now:

extension PutCommand {

    func encode() -> String {
        return "put \(self.priority) \(self.delay) \(self.timeToRun) \(body.count)\r\n\(body)\r\n"
    }

}

let command = PutCommand(priority: 10, delay: 0, timeToRun: 60, body: "download-item-62")
print(command.encode())

Now that we have a basic implementation of the put command, we can write a couple of tests.

Snapshotting

This code practically test itself: there are no dependencies, no side-effects and no mutability. You instantiate a new struct with the proper parameters and you call a method that turns it into a string. This string will later be sent to a beanstalk server. The test would roughly look like this:

final class EncodingTests: BaseTestClass {

    func testEncoding() throws {
        let command = ProducerCommand(priority: 10, delay: 0, timeToRun: 300, body: "What's up?!")
        XCTAssertEqual("put 10 0 300 11\r\nWhat's up?!\r\n", try command.encode())
    }

}

This is what you’d expect to find in any codebase relying on classic unit testing: a method returns a String and we compare it with the expected one. That said, this test was annoying to write because I had to double check that the line breaks were there.

Automated testing is supposed to bring you confidence in your code. You won’t gain much confidence if the assertions are hard to write in the first place. Plus, testing different use cases becomes really tedious if you have to write the expectations by hand.

So what’s the solution? Snapshot testing!

I decided to try that approach, after encountering Jest in Javascript, iOSSnapshotTestCase in iOS and more recently (and more assiduously), the Swift Snapshot Testing library.

I found this approach to be particularly fitting for this exercise. If you want to learn more about this library I would recommend starting with the official repository on GitHub or episode #41 of the pointfree.co series.

When using this library, the workflow slightly differs from the regular testing approach: you start by writing your assertion, using the assertSnapshot function:

final class EncodingTests: BaseTestClass {

    func testEncodingFulltextCommand() throws {
        let command = ProducerCommand(priority: 10, delay: 0, timeToRun: 300, body: "What's up?!")
        try assertSnapshot(matching: command.encode(), as: .lines)
    }

}

Then you run your test. It will fail the first time, but a snapshot will be created in a __SNAPSHOT__ directory. Finally, you check that the snapshot looks good and check it into your repository.

A snapshot generated by the library, opened in a text editor

The next time you run your test, the test will compare the result of the call to encode to the content of the saved snapshot file. If they match, the test will pass. If they don’t, you will get a diff, showing where the error is. Snapshot Testing is particularly amazing, since you will be able to test different cases (trying to encode an invalid body, a body too large, etc.) without having to come up with the assertion yourself.

In summary, when it comes to using snapshotting to test your code, the workflow is a bit different from regular unit testing: you start by writing the implementation, then you run a failing test that will save a snapshot, and finally you run it again to make it pass.

The remaining problem is making sure that the generated snapshot makes sense, and when you’re implementing a protocol it can be hard. For example, you may have missed a line break, the number of bytes might be too big, the command name might be misspelled… Only a running beanstalk server can truly validate the message.

Let’s take a step back and check how we can make sure that our snapshots are valid, even if they seem to be when viewing them in a text editor. To achieve that, we’ll send the content of our snapshot to a running beanstalk instance.

Testing with a real server

Deploying a Beanstalk server is pretty simple with docker. You can get a server up and running in seconds by running the following command in a terminal prompt:

docker run -d -p 11300:11300 schickling/beanstalkd

This will start a beanstalk server in the background (that’s the -d flag), that will be listening on port 11300 (that’s the -p flag, it maps the container port 11300 to the port 11300 on the host). You can read more about port mapping with docker on this page.

The beanstalk protocol is designed to run over TCP. You can talk to a server that uses TCP with the nc(or netcat) command line tool that’s available on most platforms, including macOS:

nc localhost 11300 < path/to/tests/__SNAPSHOTS__/testCommand.1.txt
# Will print the response from the beanstalk server: `INSERTED <id>` where id is
# is an integer describing the id of job that was sent to the server

This is one of the reasons why snapshot testing is great: not only the tests are easy to write, but the generated snapshots can also be validated in different ways, to reenforce your confidence in the implementation.

There are limitations to validating snapshots by sending them to a running server. For example: you may try to send a command to the beanstalk server to delete a job that does not exist. The response from the server will tell you that the job was not found, instead of sending you the confirmation that the job was deleted. This probably belongs in the realm of end-to-end or integration testing.