lib/simplestate/state_holder.rb in simplestate-0.1.0 vs lib/simplestate/state_holder.rb in simplestate-0.1.1

@@ -5,16 +5,17 @@
     # Set current_state within SimpleDelegator
     super(initial_state)
   end
 
   def transition_to(new_state_class)
-    current_state.send(:exit) if current_state
+    current_state.send(:exit)
     set_new_state(new_state_class)
     current_state.send(:enter)
   end
 
   def set_new_state(new_state_class)
-    current_state = new_state_class.new(self, current_state)
+    new_state = new_state_class.new(self, current_state.class)
+    self.current_state = new_state
   end
 
   def current_state
     __getobj__
   end