Refactor StateMachine

Co-authored-by: Szymon Fiedler <szymon.fiedler@gmail.com>

lib/booking/appointment.rb

@@ -5,35 +5,41 @@ class Appointment
     def initialize(id, clock:)
       @id = id
       @clock = clock
-      @state = StateMachine.new(%w[proposed accepted rejected])
+      @state =
+        StateMachine.new(
+          {
+            proposed: %i[accepted rejected cancelled],
+            accepted: [:cancelled],
+            rejected: [],
+            cancelled: []
+          },
+          method(:invalid_transition)
+        )
     end
 
     def propose
-      invalid_state unless @state.initial?
       @state.to_proposed
       AppointmentProposed.new(id: @id, proposed_at: @clock.call)
     end
 
     def accept
-      invalid_state unless @state.proposed?
       @state.to_accepted
       AppointmentAccepted.new(id: @id, accepted_at: @clock.call)
     end
 
     def reject
-      invalid_state unless @state.proposed?
       @state.to_rejected
       AppointmentRejected.new(id: @id, rejected_at: @clock.call)
     end
 
     def cancel
-      invalid_state unless @state.accepted? || @state.proposed?
+      @state.to_cancelled
       AppointmentCancelled.new(id: @id, cancelled_at: @clock.call)
     end
 
     private
 
-    def invalid_state
+    def invalid_transition
       raise Error
     end
   end

lib/booking/state_machine.rb

@@ -1,19 +1,29 @@
 module Booking
   class StateMachine
-    def initialize(states)
-      states.each do |state|
-        define_singleton_method "#{state}?" do
-          @state == state
-        end
+    INITIAL = Data.define
+
+    def initialize(states, on_error)
+      @states = states
+      @current_state = INITIAL
 
+      states.each_key do |state|
         define_singleton_method "to_#{state}" do
-          @state = state
+          on_error.call unless valid_transition?(state)
+          @current_state = state
         end
       end
     end
 
-    def initial?
-      @state.nil?
+    private
+
+    attr_reader :current_state
+
+    def valid_transition?(next_state)
+      if current_state == INITIAL
+        @states.keys.first == next_state
+      else
+        @states.fetch(current_state).include?(next_state)
+      end
     end
   end
 end

test/lib/booking/appointment_test.rb

@@ -116,6 +116,14 @@ class AppointmentTest < ActiveSupport::TestCase
       assert_raises(Appointment::Error) { appointment.cancel }
     end
 
+    test "cannot reject after cancel" do
+      appointment = new_appointment
+      appointment.propose
+      appointment.cancel
+
+      assert_raises(Appointment::Error) { appointment.reject }
+    end
+
     private
 
     def appointment_id
@@ -135,31 +143,19 @@ def new_appointment
     end
 
     def appointment_proposed
-      AppointmentProposed.new(
-        id: appointment_id,
-        proposed_at: current_time
-      )
+      AppointmentProposed.new(id: appointment_id, proposed_at: current_time)
     end
 
     def appointment_accepted
-      AppointmentAccepted.new(
-        id: appointment_id,
-        accepted_at: current_time
-      )
+      AppointmentAccepted.new(id: appointment_id, accepted_at: current_time)
     end
 
     def appointment_rejected
-      AppointmentRejected.new(
-        id: appointment_id,
-        rejected_at: current_time
-      )
+      AppointmentRejected.new(id: appointment_id, rejected_at: current_time)
     end
 
     def appointment_cancelled
-      AppointmentCancelled.new(
-        id: appointment_id,
-        cancelled_at: current_time
-      )
+      AppointmentCancelled.new(id: appointment_id, cancelled_at: current_time)
     end
   end
 end