SWE 430: Example 2 --- Refactoring

This lesson will demonstrate refactoring. We will start with an existing program and refactor it to remove duplication and improve the structure of the program with the end result of increased readability and maintainability of the code.

Problem Description

You join a project to update a sentry gun controller. The controller is modeled as a finite state machine (FSM) with four states: searching, tracking, firing, and waiting. Figure 1 shows the controller FSM. Each box represents a state. The controller receives events from its sensors and can send signals to direct the sentry gun. Events are shown as labels on the transition arrows and the signals are prefixed by 's:' in Figure 1. Each state has an associated signal. When the controller enters a state, it emits the corresponding signal. The signal is only sent when transitioning from one state to a different state and not when an event causes the controller to remain in the same state.

[States] The Controller's states:

waiting
The system is out of ammunition and waiting to be restocked.
searching
The system does not have a valid target. It is searching for a new target to attack.
tracking
The system is tracking a valid target but has overheated; it must cool down before it can fire on the target.
firing
The system is tracking and firing on a target.

[Events] Events that may be sent to the Controller:

selected
A valid target is selected (target).
dead
Current target is dead.
out_of_range
Current target is out of range
too_hot
Temperature has become too hot.
temp_ok
Temperature has dropped to an acceptable level.
out_of_ammo
Out of ammunition.
stocked
Ammunition restocked.

[Signals] The Controller may send the following signals to the sentry gun:

search
Search for a new target.
track_only
Track current target, but do not fire.
track_and_fire
Track current target and fire at will.
wait
Wait for ammunition to be restocked.

Figure 1: Finite State Machine Model of the Controller

Finite state machine model of the Controller class

Thankfully, the previous developer who designed and implemented the controller created unit tests for the system. Now the controller must be updated to handle a more complex set of situations. You have determined that the controller will need several new states, events and signals. Looking at the original code, you see that the structure of the program is not suitable to make these changes. Therefore, you decide to refactor the original code into a structure more suitable for the necessary changes.

Original Controller Program

The original controller program is listed below.

original.coffee
window.sys = {}


class sys.Controller
    constructor: (@gun) ->
        @state = 'searching'
        @temperature = 'temp_ok'
        @target = null

    set_state: (state) ->
        @state = state

    get_state: ->
        @state

    handle_event: (event) ->
        switch @state
            when 'waiting'
                switch event.name
                    when 'stocked'
                        @gun.send_signal 'search'
                        @set_state 'searching'
                    when 'too_hot', 'temp_ok'
                        @temperature = event.name
                    else
                        throw new sys.Error event.name, @state
            when 'searching'
                switch event.name
                    when 'selected'
                        @target = event.target
                        if @temperature == 'too_hot'
                            @gun.send_signal 'track_only'
                            @set_state 'tracking'
                        else if @temperature == 'temp_ok'
                            @gun.send_signal 'track_and_fire'
                            @set_state 'firing'
                        else
                            throw new sys.Error event.name, 
                                                @state,
                                                @temperature
                    when 'too_hot', 'temp_ok'
                        @temperature = event.name
                    else
                        throw new sys.Error event.name, @state
            when 'tracking'
                switch event.name
                    when 'dead', 'out_of_range'
                        @target = null
                        @gun.send_signal 'search'
                        @set_state 'searching'
                    when 'too_hot'
                        @temperature = event.name
                    when 'temp_ok'
                        @temperature = event.name
                        @gun.send_signal 'track_and_fire'
                        @set_state 'firing'
                    else
                        throw new sys.Error event.name, @state
            when 'firing'
                switch event.name
                    when 'dead', 'out_of_range'
                        @target = null
                        @gun.send_signal 'search'
                        @set_state 'searching'
                    when 'too_hot'
                        @temperature = event.name
                        @gun.send_signal 'track_only'
                        @set_state 'tracking'
                    when 'temp_ok'
                        @temperature = event.name
                    when 'out_of_ammo'
                        @target = null
                        @gun.send_signal 'wait'
                        @set_state 'waiting'
                    else
                        throw new sys.Error event.name, @state
            else
                throw new sys.Error event.name, @state


class sys.Error
    constructor: (@event, @state, @temperature=null) ->
        tail = if @temperature is null
            ''
        else
            " with temp: '#{@temperature}'"
        @message = "'#{@event}' event in '#{@state}' state" + tail
    
    toString: ->
        @message

Figure 2: Class Diagram of the Original Program

Class diagram of the original program

Figure 2 shows the class diagram of the original program. The original code also has a set of associated unit tests. We will use these unit tests as our safety-net while refacoring the original code. Here is the original test code in HTML. We copy the original production code and test code into new files named refactored.coffee and test-refactored.coffee, so we can compared the finished product with the originals at the end of the example.

Creating State Classes

Most of the Controller logic is embedded in the handle_event() function. Whenever we see one big function that tries to do everything, we know we have a design issue. Imagine trying to add several new states, events and signals to the above code. It would require adding several blocks to the outer switch statement and modifying many of the already established blocks. It would be better to isolate the states from each other and isolate the events within each state.

A first step may be to move each 'when' block of the outer switch statement into a separate function. However, since this problem could clearly benefit from the state design pattern, we will instead move each 'when' block into its own class. We will start with the 'waiting' state. Here is a copy of the relevant code---the [when 'waiting'] block from the handle_event() method of the Controller class:

refactored.coffee
class sys.Controller
    handle_event: (event) ->
        switch @state
            when 'waiting'
                switch event.name
                    when 'stocked'
                        @gun.send_signal 'search'
                        @set_state 'searching'
                    when 'too_hot', 'temp_ok'
                        @temperature = event.name
                    else
                        throw new sys.Error event.name, @state
            ... more code ...

We begin by creating a new 'WaitingState' class with a handle_event() method. Then we copy the above [when 'waiting'] block into the new handle_event() method. Be sure to "un-indent" the copied code to the left.

refactored.coffee
class WaitingState
    handle_event: (event) ->
        switch event.name
            when 'stocked'
                @gun.send_signal 'search'
                @set_state 'searching'
            when 'too_hot', 'temp_ok'
                @temperature = event.name
            else
                throw new sys.Error event.name, @state

The code, as it is above, is incorrect because it makes references to variables from the context of a 'Controller' instance. We need to fix these references to be valid from within the context of a 'WaitingState' instance. First we need to identify all the dependencies of the code block, and then ensure they are all somehow passed to the WaitingState.handle_event() method. To accomplish this, we first scan the code for the variables it accesses (remember, the @ symbol means 'this.' and 'this', in the original context, refers to the controller instance). The handle_event() method needs to know the 'event', the 'controller' and the 'gun'. The event can be different for each call to handle_event(), so we will make that a parameter of the handle_event() method. The controller and the gun always refer to the same instances, so we will make them instance variables of the class. This means we need a constructor that takes a 'Controller' instance and a 'Gun' instance as parameters and saves them as instance variables.

refactored.coffee
class WaitingState
    constructor: (@controller, @gun) ->

    handle_event: (event) ->
        switch event.name
            when 'stocked'
                @gun.send_signal 'search'
                @set_state 'searching'
            when 'too_hot', 'temp_ok'
                @temperature = event.name
            else
                throw new sys.Error event.name, @state

The above code includes the new constructor that creates 'Controller' and 'Gun' instance variables. Now we must check each line in the handle_event() method that makes use of the @ symbol and perform any necessary changes to ensure it is used in the right context. Line 7 above can stay the same because gun has become an instance variable of the WaitingState class just as it was for the Controller class. We insert 'controller' after the @ symbol in lines 8 and 10 above (lines 9 and 11 below) since the 'set_state' and 'temperature' variables still belong to the controller. Line 12 above references the "state" attribute of the controller instance. It does not make sense to leave this as an attribute of the controller, since any instance of the WaitingState class knows what state it is in already. By definition, when running code of a WaitingState instance, the Controller must be in the 'waiting' state. So instead, we will make an instance variable in the constructor, @name (line 3 below), which stores the name of the state---'waiting'---and fix line 12 above to reference '@name' instead (line 13 below).

refactored.coffee
class WaitingState
    constructor: (@controller, @gun) ->
        @name = 'waiting'

    handle_event: (event) ->
        switch event.name
            when 'stocked'
                @gun.send_signal 'search'
                @controller.set_state 'searching'
            when 'too_hot', 'temp_ok'
                @controller.temperature = event.name
            else
                throw new sys.Error event.name, @name

Our new WaitingState class in now correct. All our tests still pass, but our WaitingState class is not actually being used yet. Let's modify the Controller class to make use of our new WaitingState class.

refactored.coffee
class sys.Controller
    constructor: (@gun) ->
        @state = 'searching'
        @temperature = 'temp_ok'
        @target = null
        @waiting_state = new WaitingState this, @gun

    handle_event: (event) ->
        switch @state
            when 'waiting'
                @waiting_state.handle_event event
            when 'searching'
... more code ...

We create an instance of the WaitingState class and assign it to a new "waiting_state" instance variable (line 6). In line 11, we call the handle_event() method on the WaitingState instance. We delete all the original code under the [when 'waiting'] block so only the "@waiting_state.handle_event(event)" line remains. Running the tests shows everything still works as it did with the original code.

Notice, we did not alter the Controller class until the new WaitingState class was complete. We left the original Controller class in place until the replacement code was ready to be used. Successful refactoring is about taking small, manageable steps and quickly running our unit tests after each change to ensure our code still works.

Now we repeat the process for the other 3 states, "searching", "tracking", and "firing". We will create 3 new classes, one for each of the states. Since the process is identical for all 3, we will take a moment to spell out the steps of the process. Note that <name> refers to the name of the state (either "searching", "tracking", or "firing").

Create a new class:

  1. Copy the class structure of the "WaitingState" class.
  2. Change the new class name to "<name>State".
  3. Copy the code block from the [when "<name>"] block of Controller.handle_event() to <Name>State.handle_event()
  4. In the constructor, fix the line to read @name = "<name>"
  5. For each @ symbol:
refactored.coffee
class SearchingState
    constructor: (@controller, @gun) ->
        @name = 'searching'

    handle_event: (event) ->
        switch event.name
            when 'selected'
                @controller.target = event.target
                if @controller.temperature == 'too_hot'
                    @gun.send_signal 'track_only'
                    @controller.set_state 'tracking'
                else if @controller.temperature == 'temp_ok'
                    @gun.send_signal 'track_and_fire'
                    @controller.set_state 'firing'
                else
                    throw new sys.Error event.name, 
                                        @name,
                                        @controller.temperature
            when 'too_hot', 'temp_ok'
                @controller.temperature = event.name
            else
                throw new sys.Error event.name, @name


class TrackingState
    constructor: (@controller, @gun) ->
        @name = 'tracking'

    handle_event: (event) ->
        switch event.name
            when 'dead', 'out_of_range'
                @controller.target = null
                @gun.send_signal 'search'
                @controller.set_state 'searching'
            when 'too_hot'
                @controller.temperature = event.name
            when 'temp_ok'
                @controller.temperature = event.name
                @gun.send_signal 'track_and_fire'
                @controller.set_state 'firing'
            else
                throw new sys.Error event.name, @name


class FiringState
    constructor: (@controller, @gun) ->
        @name = 'firing'

    handle_event: (event) ->
        switch event.name
            when 'dead', 'out_of_range'
                @controller.target = null
                @gun.send_signal 'search'
                @controller.set_state 'searching'
            when 'too_hot'
                @controller.temperature = event.name
                @gun.send_signal 'track_only'
                @controller.set_state 'tracking'
            when 'temp_ok'
                @controller.temperature = event.name
            when 'out_of_ammo'
                @controller.target = null
                @gun.send_signal 'wait'
                @controller.set_state 'waiting'
            else
                throw new sys.Error event.name, @name

The above shows all three new state classes. In practice, I would only work on one at a time; but showed them here together for the sake of brevity. Now we need to fix the Controller class to make use of the new state classes. For each state, we need to do the following three-step process.

Replace the old "when" block with a call to the new class:

  1. In the Controller constructor, add the line
    @<name>_state = new <Name>State this, @gun
  2. In Controller.handle_event(), delete the code under
    when "<name>"
  3. In Controller.handle_event(), under the [when "<name>"] bolck, add the line
    @<name>_state.handle_event event
refactored.coffee
class sys.Controller
    constructor: (@gun) ->
        @state = 'searching'
        @temperature = 'temp_ok'
        @target = null
        @waiting_state = new WaitingState this, @gun
        @searching_state = new SearchingState this, @gun
        @tracking_state = new TrackingState this, @gun
        @firing_state = new FiringState this, @gun

    set_state: (state) ->
        @state = state

    get_state: ->
        @state

    handle_event: (event) ->
        switch @state
            when 'waiting'
                @waiting_state.handle_event event
            when 'searching'
                @searching_state.handle_event event
            when 'tracking'
                @tracking_state.handle_event event
            when 'firing'
                @firing_state.handle_event event
            else
                throw new sys.Error event.name, @state

All the tests still pass. We have successfully moved the code under the 'when' blocks into their own state classes. At this point, we have not improved the code much. But we have a structure in place that will allow us to make more improvements and eventually implement the state pattern.

Using Polymorphism in Place of the Switch Statement

Now that the state classes encode the state information and behavior, we can change the Controller.state variable to reference an actual state object instead of a string. Then we can get rid of the switch statement in the Controller.handle_event() method.

refactored.coffee
class sys.Controller
    constructor: (@gun) ->
        @state_map = 
            waiting: new WaitingState this, @gun
            searching: new SearchingState this, @gun
            tracking: new TrackingState this, @gun
            firing: new FiringState this, @gun
        @state = @state_map['searching']
        @temperature = 'temp_ok'
        @target = null

    set_state: (state_name) ->
        @state = @state_map[state_name]

    get_state: ->
        @state.name

    handle_event: (event) ->
        @state.handle_event event

Here's an explanation of the changes:

Lines 3-7:
In the constructor, instead of using "<name>_state" variables for the four state objects, we create a "state_map" object that maps state names to state objects.
Line 8:
The "state" variable no longer refers to a string, but to the actual current state object. So the "state" variable is now initialized by looking up the SearchingState object in the "state_map".
Lines 12-13:
To set the state, the Controller.set_state() method looks up the correct state object in the "state_map" given the state name and assigns it as the current state.
Lines 15-16:
The Controller.get_state() method returns the name of the current state.
Lines 18-19:
Finally, the Controller.handle_event() method collapses to one line---line 19. The logic is now delegated to the current state object's handle_event() method. Instead of having an ugly switch statement, we use polymorphism to dynamically dispatch the call on the correct state object.

Figure 3: Class diagram after adding state classes

Class diagram after adding state classes

Figure 3 gives a graphical depiction of the changes we have made. All tests pass except one; the "undefined state" test. The test fails because the state object is undefined. Let's take a look at the original test.

Undefined State Test

test-refactored.coffee
test 'undefined state', ->
    start_state = 'undefined_state'
    @controller.set_state start_state
    @event.name = 'stocked'
    throws (-> @controller.handle_event @event), 
           get_regex(@event.name, start_state),
           get_label(@event.name, start_state)

This is actually a bad test and we will have to rewrite it. The test assumes it is the Controller.handle_event() method's responsibility to detect illegal states, but it should actually be the Controller.set_state() method's responsibility. As long as the class interface is respected, the only way to set the Controller.state variable is via the Controller.set_state() method. So any attempts to set the current state to an undefined state (in other words, passing a state_name to set_state() that is not a key in the "state_map") should be caught in the Controller.set_state() method. We rewrite the test accordingly.

test-refactored.coffee
test 'undefined state', ->
    state_name = 'undefined_state'
    throws (-> @controller.set_state state_name),
           new RegExp("^Undefined state: \"#{state_name}\"$")

The test now asserts set_state() throws an appropriate error when the "state_name" is something other than the four valid state names (waiting, searching, tracking, or firing). The new test is correct, but now it fails because the get_state() method is not throwing an exception when it is passed an undefined state name. Let's fix the production code to throw an exception when the state name is undefined. To do this, we'll need a new Error class, because the one we have is only suitable for events. The first problem is with the name of our existing error class; it is too generic. Let's rename it to something more descriptive, like "EventError". Before we rename it in the production code, we should to fix our "Error" test in the test code. Here is the "Error" test:

test-refactored.coffee
test 'Error', ->
    f = -> throw new sys.Error 'selected', 'waiting'
    pattern = /^'selected' event in 'waiting' state$/
    throws f, pattern
    f = -> throw new sys.Error 'selected', 'searching', 'too_hot'
    pattern = new RegExp(
        "'selected' event in 'searching' state with temp: " +
        "'too_hot'"
    )
    throws f, pattern

This is a simple fix, just replace every occurrence of "Error" with "EventError". Note: you should use the global search and replace tool in your text editor or IDE to accomplish this. Changing variable/function/class names by hand is error prone. Lines 1, 2 and 5 above have been corrected below. Our old "Error" test is now the "EventError" test.

test-refactored.coffee
test 'EventError', ->
    f = -> throw new sys.EventError 'selected', 'waiting'
    pattern = /^'selected' event in 'waiting' state$/
    throws f, pattern
    f = -> throw new sys.EventError 'selected', 'searching', 'too_hot'
    pattern = new RegExp(
        "'selected' event in 'searching' state with temp: " +
        "'too_hot'"
    )
    throws f, pattern

Now the test fails because sys.EventError does not exist. So let's make the same changes in refactored.coffee, replacing every occurrence of "Error" with "EventError".

refactored.coffee
# class sys.Error
class sys.EventError
    constructor: (@event, @state, @temperature=null) ->
        tail = if @temperature is null
            ''
        else
            " with temp: '#{@temperature}'"
        @message = "'#{@event}' event in '#{@state}' state" + tail
    
    toString: ->
        @message


class sys.Controller
    handle_event: (event) ->
        ... more code ...
            else
                # throw new sys.Error event.name, @state
                throw new sys.EventError event.name, @state
        ... more code ...

The above code shows two examples of the changes, with the old lines commented out with the # symbol. There are 7 occurrences of "Error" in refactored.coffee; 1 as the class name sys.Error and six references to the class in Controller.handle_event(). Changing all 7 to "EventError" causes the "EventError" test to pass again.

With the renamed sys.EventError class, we are ready to introduce a new, more generic sys.Error class. Our sys.Error class only needs to take a message as a constructor parameter and return it when the toString() method is called on it.

refactored.coffee
class sys.Error extends Error
    constructor: (@message) ->
        @name = "sys.ControllerError"
    
    toString: ->
        @message


class sys.EventError extends sys.Error
    constructor: (@event, @state, @temperature=null) ->
        tail = if @temperature is null
            ''
        else
            " with temp: '#{@temperature}'"
        message = "'#{@event}' event in '#{@state}' state" + tail
        super message

In this step, in addition to creating the new "Error" class, we refactored the common error code into the sys.Error class. The "EventError" class inherits the behavior of the "Error" class because line 9 links "EventError" to "Error" with the "extends" keyword. All tests pass except our "undefined state" test. With our new error classes, we are ready to fix the Controller.set_state() method to make the last test pass.

refactored.coffee
class sys.Controller
    set_state: (state_name) ->
        if state_name of @state_map
            @state = @state_map[state_name]
        else
            throw new sys.Error "Undefined state: \"#{state_name}\""

To fix the Controller.set_state() method, all we need to do is check that the "state_name" parameter is a key in the "state_map" object. If it is not, we throw a sys.Error. We finally have the "undefined state" test passing again, along with all the other tests. Figure 4 shows our class diagram with the new error classes on the left.

Figure 4: Class diagram with new error classes

Class diagram after new error classes

A BaseState Class

We have a lot of duplication in the structure of our four State classes. We would like to replace each "when" block in the switch statements with actual functions, as defined in the state pattern. To name the methods that respond to each event, we will use the naming convention "on_<event_name>" where event_name is one of the events (too_hot, out_of_ammo, out_of_range, etc). To start toward this end, let's create a BaseState class to contain the duplicate code. The four concrete state classes will then be able to inherit the common logic from the BaseState class and convert the switch statements into sets of methods.

refactored.coffee
class BaseState
    handle_event: (event) ->
        if event.name in @valid_events
            @["on_#{event.name}"] event
        else
            throw new sys.EventError event.name, @name

    on_too_hot: ->
        @controller.temperature = 'too_hot'

    on_temp_ok: ->
        @controller.temperature = 'temp_ok'

    on_dead: -> @set_searching_state()

    on_out_of_range: -> @set_searching_state()

    set_searching_state: ->
        @controller.target = null
        @gun.send_signal 'search'
        @controller.set_state 'searching'

Here is an explanation of the BaseState class:

Lines 2-6
The handle_event() method will replace the switch logic in the subclasses. If the event is valid, it calls the appropriate event handler, else it throws an error. Each state subclass is required to have a "valid_events" instance variable which lists the names of the events it responds to. If an event name is not found in the list, then the event is not supported by the state subclass and an EventError should be thrown. So line 3 checks if the event name is in the list of valid events. If it is, it calls the matching "on_<event.name>" method (line 4). If the event name is not found in the "valid_events" list, an EventError is thrown (line 6).
Lines 8-12
Default event handlers, on_too_hot() and on_temp_ok(), are defined for the "too_hot" and "temp_ok" events. If a state subclass needs a different method, it can override the default methods by defining it's own version of the on_too_hot() and/or on_temp_ok() methods.
Lines 14-16
The BaseState class also defines default event handlers, on_dead() and on_out_of_range(), for the "dead" and "out_of_range" events. Only the "firing" and "tracking" states make use of these methods. The methods will not effect the other state classes as long as the other state classes do not list "dead" or "out_of_range" in their "valid_events" list.
Lines 18-21
Since the "dead" and "out_of_range" events cause the same behavior, a set_searching_state() method captures the duplicate code.

Refactoring the Concrete State Classes to Use the Base Class

Now let's refactor the concrete WaitingState class to take advantage of the new BaseState class.

refactored.coffee
class WaitingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'waiting'
        @valid_events = ['too_hot', 'temp_ok', 'stocked']

    on_stocked: ->
        @gun.send_signal 'search'
        @controller.set_state 'searching'

    handle_event: (event) ->
        switch event.name
            when 'stocked'
                @gun.send_signal 'search'
                @controller.set_state 'searching'
            when 'too_hot', 'temp_ok'
                @controller.temperature = event.name
            else
                throw new sys.Error event.name, @name
Line 1
The concrete WaitingState class now inherits from the BaseState class using the extends keyword.
Line 4
The WaitingState class defines a list of valid events it responds to.
Lines 6-8
Since the class can inherit the on_too_hot() and on_temp_ok() event handling methods from the BaseState class, only one event handler needs to be defined---on_stocked(). The body of the function is simply the code under the [when 'stocked'] block of the handle_event() switch statement (lines 12-14).

The tests all pass, but that only means we haven't made any syntax errors. The old handle_event() method of the concrete class is still being executed because it is now overriding the new handle_event() method of the BaseState class. We need to delete it so that the BaseState handle_event() method will be executed instead. We will delete the old handle_event() method now.

refactored.coffee
class WaitingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'waiting'
        @valid_events = ['too_hot', 'temp_ok', 'stocked']

    on_stocked: ->
        @gun.send_signal 'search'
        @controller.set_state 'searching'

All the tests pass. Let's generalize the steps we need to take to refactor the other three concrete state classes (SearchingState, TrackingState, and FiringState).

  1. Make the concrete class inherit the BaseState class using the extends keyword.
  2. Define the list of valid events the concrete state responds to and assign it to "@valid_events".
  3. In the handle_event() method, for each "when" block in the switch statement:
  4. Delete the old handle_event() method, it is on longer needed.
refactored.coffee
class SearchingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'searching'
        @valid_events = ['too_hot', 'temp_ok', 'selected']

    on_selected: (event) ->
        @controller.target = event.target
        if @controller.temperature == 'too_hot'
            @gun.send_signal 'track_only'
            @controller.set_state 'tracking'
        else if @controller.temperature == 'temp_ok'
            @gun.send_signal 'track_and_fire'
            @controller.set_state 'firing'
        else
            temp = @controller.temperature
            throw new sys.EventError event.name, @name, temp


class TrackingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'tracking'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 'out_of_range']

    on_temp_ok: ->
        @controller.temperature = 'temp_ok'
        @gun.send_signal 'track_and_fire'
        @controller.set_state 'firing'


class FiringState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'firing'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 
                         'out_of_range', 'out_of_ammo']

    on_too_hot: ->
        @controller.temperature = 'too_hot'
        @gun.send_signal 'track_only'
        @controller.set_state 'tracking'

    on_out_of_ammo: ->
        @controller.target = null
        @gun.send_signal 'wait'
        @controller.set_state 'waiting'

All four concrete state classes now extend the BaseState class and define their own "valid_events" list. Our code is now free of switch statements. Also, note:

Lines 24-27
The TrackingState class overides the default on_temp_ok() method of BaseState with its own on_temp_ok() method.
Lines 36-39
The FiringState class overrides the default on_too_hot() method of BaseState with its own on_too_hot() method.

All tests pass. We have successfully refactored some of the common code of the state classes into a BaseState class. Our latest class diagram is presented in Figure 5.

Figure 5: Class diagram with BaseState class

Class diagram with BaseState class

Moving the Target Variable into the TrackingState and FiringState Classes

Examining the target variable, we see it is only used in the "TrackingState" and "FiringState" classes. If they are the classes using the variable, they should be responsible for it, not the Controller class. Before we remove the target variable from the Controller class, we implement the target variables for the "TrackingState" and "FiringState". First we need to figure out how the "target" will be set in the "TrackingState" and "FiringState" instances. Objects of both classes should always have a valid target when they are the current state of the controller. Therefore, it makes sense that the target should be set as soon as the Controller enters a "tracking" or "firing" state. So, the code to set the target should go in the Controller.set_state() method.

refactored.coffee
class sys.Controller
    set_state: (state_name, target=null) ->
        if state_name of @state_map
            @state = @state_map[state_name]
        else
            throw new sys.Error "Undefined state: \"#{state_name}\""
        if state_name in ['tracking', 'firing']
            @state.target = target
Line 2
We add a second parameter to set_state(): 'target'. The 'target' parameter is an optional parameter with a default value of 'null'. It is only accessed if transitioning to the 'tracking' or 'firing' states.
Lines 7-8
If the new state is 'tracking' or 'firing', the states' target variable should be set to the value of the new target parameter.

Now we must check each call to set_state() in the production code and ensure that when transitioning into the "tracking" or "firing" states, the current target is passed as a second parameter to the controller.set_state() method.

refactored.coffee
class SearchingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'searching'
        @valid_events = ['too_hot', 'temp_ok', 'selected']

    on_selected: (event) ->
        @controller.target = event.target
        if @controller.temperature == 'too_hot'
            @gun.send_signal 'track_only'
            @controller.set_state 'tracking', event.target
        else if @controller.temperature == 'temp_ok'
            @gun.send_signal 'track_and_fire'
            @controller.set_state 'firing', event.target
        else
            temp = @controller.temperature
            throw new sys.EventError event.name, @name, temp

The SearchingState.on_selected() method calls the controller.set_state() method on both lines 10 and 13. We must fix both of these calls because they concern the 'tracking' and 'firing' states. Therefore, we add event.target as the second parameter for each of the calls to controller.set_state().

refactored.coffee
class TrackingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'tracking'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 'out_of_range']

    on_temp_ok: ->
        @controller.temperature = 'temp_ok'
        @gun.send_signal 'track_and_fire'
        @controller.set_state 'firing', @target


class FiringState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'firing'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 
                         'out_of_range', 'out_of_ammo']

    on_too_hot: ->
        @controller.temperature = 'too_hot'
        @gun.send_signal 'track_only'
        @controller.set_state 'tracking', @target

    on_out_of_ammo: ->
        @controller.target = null
        @gun.send_signal 'wait'
        @controller.set_state 'waiting'

Both the TrackingState.on_temp_ok() and the FiringState.on_too_hot() methods call the controller.set_state() method (occurs on lines 9 and 21). We must also fix these calls because they again concern the 'tracking' and 'firing' states. This time, we add @target as the second parameter for each of the calls to controller.set_state(). We use @target because the controller is already in a 'firing' or 'tracking' state, therefore the state object knows what the current target is and can pass that information along to the Controller.set_state() method.

All our tests pass and we have implemented all the necessary code for the 'TrackingState' and 'FiringState' classes to have a 'target' variable. Now we can remove the 'target' variable from the 'Controller' class since the other two classes have assumed the responsibility.

refactored.coffee
class BaseState
    ... more code ...

    set_searching_state: ->
        # @controller.target = null
        @gun.send_signal 'search'
        @controller.set_state 'searching'


class SearchingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'searching'
        @valid_events = ['too_hot', 'temp_ok', 'selected']

    on_selected: (event) ->
        # @controller.target = event.target
        if @controller.temperature == 'too_hot'
            @gun.send_signal 'track_only'
            @controller.set_state 'tracking', event.target
        else if @controller.temperature == 'temp_ok'
            @gun.send_signal 'track_and_fire'
            @controller.set_state 'firing', event.target
        else
            temp = @controller.temperature
            throw new sys.EventError event.name, @name, temp


class FiringState extends BaseState
    ... more code ...

    on_out_of_ammo: ->
        # @controller.target = null
        @gun.send_signal 'wait'
        @controller.set_state 'waiting'


class sys.Controller
    constructor: (@gun) ->
        @state_map = 
            waiting: new WaitingState this, @gun
            searching: new SearchingState this, @gun
            tracking: new TrackingState this, @gun
            firing: new FiringState this, @gun
        @state = @state_map['searching']
        @temperature = 'temp_ok'
        # @target = null

    ... more code ...

Lines 5, 16, 32, and 46---commented with the # symbol for emphases---refer to the obsolete controller.target variable. The lines are no longer needed and can now be deleted.

All tests now pass except one, the "is target set properly?" test. Let's take a look at the failing test.

test-refactored.coffee
test 'is target set properly?', ->
    check_target = ((controller, event) ->
        (start_state, event_name, temp, expect_null) ->
            controller.set_state start_state
            event.name = event_name
            is_null = start_state in ['waiting', 'searching']
            controller.target = if is_null then null else event.target
            controller.temperature = temp
            controller.handle_event event
            equal (controller.target is null), expect_null
    ) @controller, @event
    for temp in ['too_hot', 'temp_ok']
        check_target 'searching', 'selected', temp, false
    # Pairs that transition into searching or waiting states
    null_pairs = [ ... more code ... ]
    for [start_state, event_name] in null_pairs
        check_target start_state, event_name, 'temp_ok', true
    # Pairs that transition into tracking or firing states
    not_null_pairs = [ ... more code ... ]
    for [start_state, event_name] in not_null_pairs
        check_target start_state, event_name, 'temp_ok', false

There are two problems with this test. 1) It still expects the target variables to be attached to the controller object instead of the controller.state object. 2) It thinks the target will be null when in the 'waiting' or 'searching' states when it should be undefined instead.

Most of the problem lies in line 10. We fix this by changing line 10 above to line 8 below. Instead of referencing controller.target, we reference controller.state.target. And instead of comparing the target to "null" we compare it to "undefined". We also cleaned up the preceding lines, lines 4-9 above to yield lines 4-7 below.

test-refactored.coffee
test 'is target set properly?', ->
    check_target = ((controller, event) ->
        (start_state, event_name, temp, expect_undef) ->
            event.name = event_name
            controller.set_state start_state, event.target
            controller.temperature = temp
            controller.handle_event event
            equal (controller.state.target is undefined), expect_undef
    ) @controller, @event
    for temp in ['too_hot', 'temp_ok']
        check_target 'searching', 'selected', temp, false
    # Pairs that transition into searching or waiting states
    undef_pairs = [ ... more code ... ]
    for [start_state, event_name] in undef_pairs
        check_target start_state, event_name, 'temp_ok', true
    # Pairs that transition into tracking or firing states
    not_undef_pairs = [ ... more code ... ]
    for [start_state, event_name] in not_undef_pairs
        check_target start_state, event_name, 'temp_ok', false

To complete the changes, we renamed variables that used "null" in their name, to use "undef" instead (specifically, the variables "expect_undef", "undef_pairs", and "not_undef_pairs" in lines 3, 8, 13, 14, 17, and 18).

All tests pass once again and we have successfully moved the target variable responsibility from the Controller class into the TrackingState and FiringState classes.

Two More Simple Fixes

The Controller class is still holding a reference to the 'gun' object, but the Controller never uses it. Only the state classes need to know the 'gun' object. So we can remove the @ symbol from in front of the 'gun' variable everywhere it occurs in the Controller constructor method (5 occurrences). The constructor now looks like this:

refactored.coffee
class sys.Controller
    constructor: (gun) ->
        @state_map = 
            waiting: new WaitingState this, gun
            searching: new SearchingState this, gun
            tracking: new TrackingState this, gun
            firing: new FiringState this, gun
        @state = @state_map['searching']
        @temperature = 'temp_ok'

The SearchingState.on_selected() method has a useless else clause. It is checking for a condition that is already dealt with in the BaseState.handle_event() method. So we can safely delete the unneeded else clause. The lines to be deleted, lines 13-15, are commented-out below.

refactored.coffee
class SearchingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'searching'
        @valid_events = ['too_hot', 'temp_ok', 'selected']

    on_selected: (event) ->
        if @controller.temperature == 'too_hot'
            @gun.send_signal 'track_only'
            @controller.set_state 'tracking', event.target
        else if @controller.temperature == 'temp_ok'
            @gun.send_signal 'track_and_fire'
            @controller.set_state 'firing', event.target
        # else
        #    temp = @controller.temperature
        #    throw new sys.EventError event.name, @name, temp

Remove Code Duplication in Event Handlers

The event handling code often needs to perform two actions. 1) Send a signal to the gun (@gun.send_signal()). 2) Change the state of the controller (@controller.set_state()). Let's localize this code into a new method in the BaseState class, this way all concrete state classes can share the same code. We will call the method set_state().

refactored.coffee
class BaseState
    set_state: (signal, state_name, target=null) ->
        @gun.send_signal signal
        @controller.set_state state_name, target

    set_searching_state: ->
        # @gun.send_signal 'search'
        # @controller.set_state 'searching'
        @set_state 'search', 'searching'

    ... more code ...
Lines 2-4
We add the new set_sate() method. It takes as input a signal to send to the gun and the name of the state to transition to. It also takes an optional third argument, target, which defaults to null if not provided.
Lines 7-8
We delete lines 7 and 8 (shown commented-out) because this functionality will now be provided by the BaseState.set_state() method.
Line 9
This line simply calls set_state() with the appropriate parameters to replace the 2 deleted lines.

All tests pass. With the addition of the more general set_state() method, the set_searching_state() method is pretty useless. Let's in-line the behavior of the set_searching_state() back into the methods that use it (on_dead() and on_out_of_range()) so we can safely delete the set_searching_state() method.

refactored.coffee
class BaseState
    on_dead: -> @set_state 'search', 'searching'

    on_out_of_range: -> @set_state 'search', 'searching'

    set_state: (signal, state_name, target=null) ->
        @gun.send_signal signal
        @controller.set_state state_name, target
    #
    # set_searching_state: ->
    #     @set_state 'search', 'searching'

Lines 2 and 4 above show the on_dead() and on_out_of_range() methods making use of the new set_state() method instead of the set_searching_state() method. Lines 9-11 can now be deleted. All test pass. We can convert the rest of the event handlers to use the new BaseState.set_state() method where appropriate.

refactored.coffee
class WaitingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'waiting'
        @valid_events = ['too_hot', 'temp_ok', 'stocked']

    on_stocked: ->
        # @gun.send_signal 'search'
        # @controller.set_state 'searching'
        @set_state 'search', 'searching'


class SearchingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'searching'
        @valid_events = ['too_hot', 'temp_ok', 'selected']

    on_selected: (event) ->
        if @controller.temperature == 'too_hot'
            # @gun.send_signal 'track_only'
            # @controller.set_state 'tracking', event.target
            @set_state 'track_only', 'tracking', event.target
        else if @controller.temperature == 'temp_ok'
            # @gun.send_signal 'track_and_fire'
            # @controller.set_state 'firing', event.target
            @set_state 'track_and_fire', 'firing', event.target


class TrackingState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'tracking'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 'out_of_range']

    on_temp_ok: ->
        @controller.temperature = 'temp_ok'
        # @gun.send_signal 'track_and_fire'
        # @controller.set_state 'firing', @target
        @set_state 'track_and_fire', 'firing', @target


class FiringState extends BaseState
    constructor: (@controller, @gun) ->
        @name = 'firing'
        @valid_events = ['too_hot', 'temp_ok', 'dead', 
                         'out_of_range', 'out_of_ammo']

    on_too_hot: ->
        @controller.temperature = 'too_hot'
        # @gun.send_signal 'track_only'
        # @controller.set_state 'tracking', @target
        @set_state 'track_only', 'tracking', @target

    on_out_of_ammo: ->
        # @gun.send_signal 'wait'
        # @controller.set_state 'waiting'
        @set_state 'wait', 'waiting'

There are 6 areas that can make use of the new set_state() method: lines 7-9, 19-21, 23-25, 35-37, 48-50, and 53-55. In each case, we delete the 2 lines that individually call gun.send_signal() and controller.set_state() and replace them with a single call to set_state with the appropriate parameters. With the new changes, all of the tests pass and we have completed our refactoring.

Figure 6: Class diagram of completed refactoring

Class diagram of completed refactoring

By making all these changes, we have isolated the states from each other and isolated the events within each state. The new structure of the code makes it far easier to add new states, events and signals to the controller program. Figure 6 depicts the class diagram of the final program.

Here is the original code, with its associated original test code and the refactored code, with its associated refactored test code for your convenience.

Although the refactored code is a few lines longer than the original, it is actually less code. The original.coffee is 3.4 KB while the refactored.coffee is 3.2 KB. The refactored code is much less dense than the original.coffee. This is due to the larger number of blank lines in the refactored.coffee that are used to separate the larger number of classes and methods from each other. Often, refactoring will lead to a greater decreases in program size. In any case, refactoring should always improve the structure and design of the software.

Now you are ready to tackle exercise 2.

Lyall Jonathan Di Trapani 08 Mar 2013