lib/hat_trick/wizard_steps.rb in hat-trick-0.0.1 vs lib/hat_trick/wizard_steps.rb in hat-trick-0.1.0

@@ -24,21 +24,24 @@
     end
 
     def steps_after(_step)
       step = find_step(_step)
       return [] unless step
-      after_index = steps.index(step) + 1
+      step_index = steps.index(step)
+      max_index = steps.count - 1
+      after_index = step_index < max_index ? step_index + 1 : step_index
       steps[after_index .. -1]
     end
 
     def step_before(_step)
       steps_before(_step).last
     end
 
     def steps_before(_step)
       step = find_step(_step)
       return [] unless step
-      before_index = steps.index(step) - 1
+      step_index = steps.index(step)
+      before_index = step_index > 0 ? step_index - 1 : step_index
       steps[0 .. before_index]
     end
 
     def add_step(step, args={})
       if step.is_a?(HatTrick::Step) || step.is_a?(HatTrick::StepDefinition)