Encapsulating More than Just Resources with Test Context Classes

In my last article, I explored automatic resource management in Scala by developing test context classes to reliably manage resources in functional tests. I have explored the topic of test context classes a few times, and each time I focused on using these classes purely for resource management. In this article, I will examine extending these classes to include more than just resource management and show how they can make functional test suites easier to write, more reliable, and easier to maintain.

In my last article, I defined the following UsingWebSocket object.

object UsingWebSocket {
  def apply[T](f: WebSocketClient => T)
              (implicit url: String, callback: String => Unit): T = {
    val client = WebSocketClient(url, callback)
    assert(client.connectBlocking(), s"Failed to connect to $url")
    try {
      f(client)
    } finally {
      client.close()
    }
  }
}

I used it in this simple test of a client sending a message to a WebSocket server.

class WebSocketTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit def callback(response: String): Unit = {}

  "Client" should "send message to server" in {
    UsingWebSocket { client =>
      client.send("Bueller?")
    }
  }
}

The example of testing a WebSocket server is a perfect situation for extending test context classes to streamline functional tests. Unlike HTTP, which is a request-response protocol, the WebSocket protocol is full-duplex, meaning the client and the server can send messages to each other independently, even simultaneously.

Consider a WebSocket server designed for an Internet of Things (IoT) application, where the client can subscribe to various topics in a typical publish-subscribe model. When an event occurs for a topic that the client has subscribed to, the server will notify the client. To subscribe to a topic, the client sends a JSON subscription message to the server and includes a unique subscription identifier. The server acknowledges the subscription request with an acknowledgment message, also in JSON format, and includes the unique subscription identifier assigned by the client. Here is what a functional test for this subscription mechanism might look like, using the UsingWebSocket test context class to manage the WebSocket connection.

class SubscriptionAcknowledgmentTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit val timeout = 5 seconds

  "Server" should "acknowledge a subscription request" in {
    val id = uuid()
    var success = false

    implicit def callback(response: String): Unit = {
      response should be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1"]}""")
      success = true
    }

    UsingWebSocket { client =>
      val request = s"""{"message":"subscribe","id":"$id","subscriptions":["/channels/1"]}"""
      client.send(request)
      Thread.sleep(5000)
    }

    assert(success, "No subscription acknowledgement from server")
  }
}

This test is complicated by the fact that the WebSocket protocol is not a request-response protocol. The test must block for a few seconds to allow the response to be received in the callback method and it requires signaling between the callback method and the test method proper, via the success Boolean, to determine the outcome of the test. This pattern isn't horrible, but the intent of the test is not immediately clear, especially given that there are two assertions in the test.

Let's also consider that the server sends status and control messages to the client. There is a chance that the first message received by the client will be a status or control message, rather than the subscription acknowledgment. The test needs to filter out these messages and only consider the subscription acknowledgment message.

class SubscriptionAcknowledgmentTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit val timeout = 5 seconds

  "Server" should "acknowledge a subscription request" in {
    val id = uuid()
    var success = false

    implicit def callback(response: String): Unit = {
      if (response.startsWith( """{"message":"ack"""")) {
        response should
          be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1"]}""")
        success = true
      }
    }

    UsingWebSocket { client =>
      val request = s"""{"message":"subscribe","id":"$id","subscriptions":["/channels/1"]}"""
      client.send(request)
      Thread.sleep(5000)
    }

    assert(success, "No subscription acknowledgement from server")
  }
}

This test is still not horrible, but the intent of this test is significantly obscured. Because the specific test assertion is part of the callback, this code will need to be repeated for each test in this test suite. To simplify this, we could encapsulate the callback in another class that takes just the unique test assertion as a function literal, but I think this is a case where extending the UsingWebSocket context class to perform more than just resource management provides a cleaner solution.

I'll start by defining a WebSocketResponse class that uses the UsingWebSocket class, but awaits the response with a timeout.

object WebSocketResponse {
  def apply(message: String)(predicate: String => Boolean)
           (implicit timeout: Duration, url: String): String = {
    val promise = Promise[String]()
    val future = promise.future

    implicit def callback(response: String): Unit = {
      if (predicate(response)) {
        promise.success(response)
      }
    }

    UsingWebSocket { client =>
      client.send(message)
      Await.result(future, timeout)
    }
  }
}

This class takes a predicate to allow filtering specific messages. It is used by the following SubscriptionAcknowledgement class to send a subscription message to the server and return the subsequent subscription-acknowledgment message.

object SubscriptionAcknowledgment {
  def apply(message: String)(implicit timeout: Duration, url: String): String = {
    WebSocketResponse(message) {
      _.startsWith("""{"message":"ack"""")
    }
  }
}

The SubscriptionAcknowledgment class can then be used to write a very streamlined test with a clear assertion. The test can be written in a straightforward request-response pattern, even though the underlying system is not request-response.

class SubscriptionAcknowledgmentTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit val timeout = 5 seconds

  "Server" should "acknowledge a subscription request" in {
    val id = uuid()
    val request = s"""{"message":"subscribe","id":"$id","subscriptions":["/channels/1"]}"""
    val response = SubscriptionAcknowledgment(request)
    response should be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1"]}""")
  }
}

It is easy to define other classes that also use the WebSocketResponse class to write tests for additional message types, like the status and the control messages mentioned earlier. Note too that since the WebSocketResponse class awaits the response message and returns as soon as it is received, the test reliably completes in just a few milliseconds, rather than blocking for a fixed amount of time like the preceding tests. This means that the test suite will execute much faster and it will be less prone to timing issues.

With the SubscriptionAcknowledgment class, it is easy to write a whole suite of tests for the subscription acknowledge mechanism. The tests are reliable and have clear assertions.

class SubscriptionAcknowledgmentTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit val timeout = 5 seconds

  "Server" should "acknowledge a subscription request" in {
    val id = uuid()
    val request = s"""{"message":"subscribe","id":"$id","subscriptions":["/channels/1"]}"""
    val response = SubscriptionAcknowledgment(request)
    response should be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1"]}""")
  }

  it should "acknowledge multiple subscription requests" in {
    val id = uuid()
    val request =
      s"""{"message":"subscribe","id":"$id","subscriptions":["/channels/1","/channels/2"]}"""
    val response = SubscriptionAcknowledgment(request)
    response should
      be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1","/channels/2"]}""")
  }

  it should "not acknowledge an invalid subscription request" in {
    val id = uuid()
    val request = s"""{"message":"subscribe","id":"$id","subscriptions":["/invalid/1"]}"""
    val response = SubscriptionAcknowledgment(request)
    response should be(s"""{"message":"ack","id":"$id","subscriptions":[]}""")
  }

  it should "only acknowledge valid subscriptions for multiple subscription requests" in {
    val id = uuid()
    val request =
      s"""{"message":"subscribe","id":"$id","subscriptions":["/invalid/1","/channels/1"]}"""
    val response = SubscriptionAcknowledgment(request)
    response should
      be(s"""{"message":"ack","id":"$id","subscriptions":["/channels/1"]}""")
  }

  // And so on...
}

Functional tests like this are always reasonably involved. You'll find that as different people write these tests, they'll each have their own style for writing them. Some people will anticipate things like needing to filter out the status and control messages, whereas others will not and this filtering will be tacked on piecemeal in order to address sporadic test failures. Sleep statements will be added, adjusted, and readjusted to address timing issues. As test suites evolve in this manner, the tests become complex to understand, troubleshoot, and update. Diagnosing test failures becomes as much about understanding what the test is attempting to accomplish, as determining the reason for the test failure.

Extending test context classes to include complex and repeated operations, in addition to resource management, provides consistency and reliability across tests. It means that the tester can focus on writing the specific test, rather than on all of the orchestration required to setup the test and obtain a result that can be used for testing assertions. This consistency really helps when troubleshooting test failures, especially for tests that you did not write yourself. Test context classes also make the tests concise. The intent of each test is clear, because the test assertions are clear and they are not lost in the test orchestration. It also means the tests are easier to change as the system evolves, because the most complex code in the test is encapsulated in the test context class, rather than the test itself.