lib/hat_trick/wizard_steps.rb in hat-trick-0.1.2 vs lib/hat_trick/wizard_steps.rb in hat-trick-0.1.3

@@ -1,6 +1,8 @@
 module HatTrick
+  class StepNotFound < StandardError; end
+
   module WizardSteps
     include Enumerable
     attr_reader :steps
 
     delegate :each, :empty?, :first, :last, :to => :steps
@@ -24,13 +26,15 @@
     end
 
     def steps_after(_step)
       step = find_step(_step)
       return [] unless step
+      return [] if step.last? && !step.skipped?
       step_index = steps.index(step)
       max_index = steps.count - 1
-      after_index = step_index < max_index ? step_index + 1 : step_index
+      return [] if step_index >= max_index
+      after_index = step_index + 1
       steps[after_index .. -1]
     end
 
     def step_before(_step)
       steps_before(_step).last
@@ -38,19 +42,24 @@
 
     def steps_before(_step)
       step = find_step(_step)
       return [] unless step
       step_index = steps.index(step)
-      before_index = step_index > 0 ? step_index - 1 : step_index
+      return [] if step_index <= 0
+      before_index = step_index - 1
       steps[0 .. before_index]
     end
 
     def add_step(step, args={})
       if step.is_a?(HatTrick::Step) || step.is_a?(HatTrick::StepDefinition)
         new_step = step
       else
-        step_args = args.merge(:name => step)
+        step_args = args.merge(:name => step, :wizard => self)
         new_step = HatTrick::StepDefinition.new(step_args)
+      end
+
+      if steps.count == 0
+        new_step.first = true
       end
 
       steps << new_step
 
       new_step