Avoiding bugs in Ruby code using the state pattern

“Be more careful” is often not useful advice. Software bugs are inevitable, but as I become proficient in software development, my code tends to have fewer initial bugs.

In some part, this is because experience has taught me to spot bugs more quickly. But more importantly, I’ve learned techniques to structure code in ways that make bugs much less likely to arise in the first place.

One such technique is the state pattern. In this article, I’ll take you through an example, at first with a brittle implementation, and then reworked using the state pattern into something more stable.

The example: A messaging platform

The example I’ll use throughout this article is a simplified messaging platform. On this platform, any person can register:

account = Account.register(account_details)

The account details, passed to the .register method, would include at least an email address, but also properties such as first name and last name.

Before the registered account is usable, the person needs to confirm the account, using a code they received via email:

account.confirm("ABC123")

Then, an administrator is expected to approve the account before it can be used:

account.approve

With the account confirmed (by the person registering) and approved (by an administrator), the person is now free to use their account to send messages to other people:

account.message(
  tim,
  "I’m selling these fine leather jackets.",
)

It is important that people cannot send messages unless their account is both confirmed and approved.¹¹1 That should cut down on the cryptocurrency/NFT spam. That is unfortunately still a thing, right?

Let’s take a look at two different ways of implementing this: a naïve implementation, and a safer implementation that uses the state pattern.

A naïve implementation

The simplest implementation (that I can think of) starts with an initializer:

class Account
  def initialize
    @code = "ABC123"
  end

The initializer sets the confirmation code. In a real-world scenario, the confirmation code would be randomly generated. Here, it is hardcoded for the sake of simplicity.

We’ll also need an implementation for the Account.register method, which we used earlier:

  def self.register(account_details)
    new
  end

For this demo implementation, all it does is create a new instance of Account. In a real-world scenario, this might be the place to create a new database record and send out an email with the confirmation code. For simplicity, all of that is left out.

We’ll need #confirm, which checks whether the given code is correct, and advances the state to :confirmed:

  def confirm(code)
    if code != @code
      raise InvalidConfirmationCode
    end

    @state = :confirmed
  end

Next up is #approve, which advances the state to :confirmed_and_approved:

  def approve
    @state = :confirmed_and_approved
  end

Finally, we have #message, which is the method used for sending messages to other accounts. It requires that the account is in the :confirmed_and_approved state:

  def message(who, text)
    unless @state == :confirmed_and_approved
      raise AccountNotApproved
    end

    puts "Sending message to #{who}"
  end
end

In this example implementation of #message, we just log the message.

We’ll also need these two exception classes:

AccountNotApproved = Class.new(StandardError)
InvalidConfirmationCode = Class.new(StandardError)

Here is an example that ties it all together:

tims_account_id = 12358

account = Account.new
account.confirm("ABC123")
account.approve
account.message(tims_account_id, "What’s up, Tim?")

Here, we confirm the account, then an administrator approves the account, and lastly we use the account to send a message to Tim, whose account ID is 12358. The terminal output now shows this:

Sending message to 12358.

So far, so good.

But what if we confirm the account again, after it has already been approved?

tims_account_id = 12358

account = Account.new
account.confirm("ABC123")
account.approve
account.confirm("ABC123")
account.message(tims_account_id, "What’s up, Tim?")

Running this example results in an error during the call to the #message method:

account.rb:9:in `message':
  AccountNotApproved (AccountNotApproved)

AccountNotApproved?! The account definitely was approved. The issue is that our implementation of #confirm set the state to :confirmed:

  def confirm(code)
    if code != @code
      raise InvalidConfirmationCode
    end

    @state = :confirmed
  end

Here, we have a bug: the state shouldn’t be set to :confirmed if the account is already approved. One quick fix for this would be to zero out the @code variable, so that attempting to confirm again would fail:

  def confirm(code)
    if code != @code
      raise InvalidConfirmationCode
    end

    @state = :confirmed
    @code = ""
  end

That is a bit of a hack, though. A slightly better way to solve this is to only allow confirmation when the @state is the initial state:

  def confirm(code)
    return if @state != :initial

    if code != @code
      raise InvalidConfirmationCode
    end

    @state = :confirmed
  end

We’d also need to set the initial @state in the initializer:

  def initialize
    @code = "ABC123"
    @state = :initial
  end

With that bug solved, let us take a look at another issue: it is possible for an account to be approved without having gone through confirmation at all:

tims_account_id = 12358

account = Account.new
account.approve
account.message(tims_account_id, "What’s up, Tim?")

The terminal output now shows this:

Sending message to 12358.

This might be a bug, or it might not be. Perhaps it is intentional that administrators can approve accounts, skipping confirmation. Or perhaps this is an oversight in our implementation.

If we assume it is an oversight, and thus a bug, one way of fixing it is to verify that the state is what we expect it to be:

  def approve
    raise AccountNotConfirmed if @state != :confirmed

    @state = :confirmed_and_approved
  end

We’ll also need to define a new exception:

AccountNotConfirmed = Class.new(StandardError)
AccountNotApproved = Class.new(StandardError)
InvalidConfirmationCode = Class.new(StandardError)

The code is now — as far as I can tell — bug-free, though I’m not comfortable with the end result. This code feels brittle to me: any change in the future has a high chance of inadvertently modifying the expected behavior.

This could would need an extensive test suite. Such a test suite would verify that going through all the different paths, in all different orders, yield correct results. The presence of a test suite would make me feel more comfortable, but there is more that we can do.

Let’s now take a look at a new implementation, which uses @state rather differently.

A safer implementation

In this implementation, the Account class no longer contains the functionality for confirming, approving, and messaging. Rather, all of that is delegated to @state, which is now its own object:

class Account
  attr_accessor :state

  def initialize
    @state = InitialAccountState.new(self)
  end

  def confirm(code)
    @state.confirm(code)
  end

  def approve
    @state.approve
  end

  def message(who, text)
    @state.message(who, text)
  end
end

You could go fancy and use Ruby’s built-in Forwardable module for that if you want. With that module, the implementation of Account could look like this — doing exactly the same:

class Account
  extend Forwardable
  def_delegators :@state, :confirm, :approve, :message

  attr_accessor :state

  def initialize
    @state = InitialAccountState.new(self)
  end
end

Here is what InitialAccountState looks like:

class InitialAccountState
  def initialize(account)
    @code = "ABC123"
    @account = account
  end

  def confirm(code)
    if code != @code
      raise InvalidConfirmationCode
    end

    @account.state = ConfirmedAccountState.new(@account)
  end
end

The InitialAccountState#confirm method is quite similar to the original #confirm: it checks the confirmation codes, and raises an exception if they don’t match.

While the original #confirm changed the state using @state = :confirmed, this new #confirm method changes the account state to a new state object. (We’ll get to the implementation of ConfirmedAccountState in a bit.)

Also worth noting is that the confirmation code lives in the InitialAccountState instance. That is the only place where it is useful. It could also live in Account, but it wouldn’t have a purpose there.

Without having the other state objects implemented, we can already see that one of the buggy behaviors from earlier no longer silently succeeds:

tims_account_id = 12358

account = Account.new
account.approve
account.message(tims_account_id, "What’s up, Tim?")

This piece of code raises a NoMethodError:

account_with_state.rb:13:
  in `approve': undefined method `approve' for
  #<InitialAccountState:0x0000000100907fa8 …>
  (NoMethodError)

  @state.approve(code)
        ^^^^^^^^
      from account_with_state.rb:66:in `<main>'

This NoMethodError is intentional! It signals that there hasn’t been any thought put into the situation where someone would try to approve an account that hasn’t been confirmed yet.

This NoMethodError is a replacement for undefined behavior. I believe that it is preferable to get a NoMethodError than to execute incorrect behavior.

We are still able to implement #approve here, if we wish. Perhaps it is desirable to approve accounts from their initial state, bypassing confirmation. If so, we could implement #approve as follows:

class InitialAccountState
  …

  def approve
    @account.state =
      ConfirmedAndApprovedAccountState.new(@account)
  end
end

Let’s move on to ConfirmedAccountState:

class ConfirmedAccountState
  def initialize(account)
    @account = account
  end

  def approve
    @account.state =
      ConfirmedAndApprovedAccountState.new(@account)
  end
end

The #approve method exists here, and moves the state forward to the “account approved” state.

The ConfirmedAccountState state has no #confirm method here. If we were to try to confirm an already approved account, we’d get a NoMethodError:

account_with_state.rb:28:
  in `confirm': undefined method `confirm' for
  #<ConfirmedAccountState:0x0000000100907fa8 …>
  (NoMethodError)

  @state.confirm(code)
        ^^^^^^^^
      from account_with_state.rb:66:in `<main>'

Here, the NoMethodError is less desirable. I would probably implement #confirm anyway, and have it do nothing:

class ConfirmedAccountState
  …

  def confirm
    # Already confirmed; do nothing
  end

This makes the behavior explicit: confirming an account that is already confirmed does nothing. The #confirm method in our naïve implementation also did nothing in this case, but it wasn’t nearly as explicit.

Lastly, we have ConfirmedAndApprovedAccountState:

class ConfirmedAndApprovedAccountState
  def initialize(account)
    @account = account
  end

  def message(who, text)
    puts "Sending message to #{who}"
  end
end

This state is the least interesting: it just has the #message for sending messages.

For this state, it makes sense to implement the #confirm and #approve methods, and have them do nothing:

class ConfirmedAndApprovedAccountState
  …

  def confirm
    # Already confirmed; do nothing
  end

  def approve
    # Already approved; do nothing
  end

With all that, we have an implementation where the state transitions are explicit, and it is also more clear what actions can be taken in which states.

Closing thoughts

The implementation which uses the state pattern makes me the most comfortable. It avoids undefined behavior, and even though NoMethodErrors are a little nasty, they’re better to have than undefined behavior.

There is nothing preventing us from implementing all methods, and avoiding NoMethodError altogether. For each combination of state and method, we’d have to think about what the behavior should be. Perhaps it is doing nothing, perhaps it is raising a specific exception, or perhaps it is doing something else entirely.

Defining the behavior for each combination of state and method is not always easy. This can make the state pattern look cumbersome and difficult. However, if we want high-quality software, we can’t get away from defining behavior anyway. It seemed to be optional in the original, naïve implementation, but that led to bugs as a result.

I prefer explicit, readable code over compact code every time. Explicit code makes it easier to find and prevent bugs. Explicit, readable code is less likely to break over time, even in a codebase with a large amount of churn.

Back in university, a professor said to me that if statements are a code smell. While I think that is an over-generalisation, you’ll find that the original, naïve implementation has quite some ifs — especially in the bug fixes — while the state pattern implementation has few.

If in the future I find myself in a situation where behavior depends on state, you can be sure I’ll whip out the state pattern.