lib/fluent/factory.rb in fluent-0.6.0 vs lib/fluent/factory.rb in fluent-0.7.0

@@ -1,18 +1,25 @@
+require 'fluent/workflows'
+
 module Fluent
   module Factory
-
+    include Workflow
+    
     # Creates a definition context for actions. If an existing context
     # exists, that context will be re-used.
     #
     # @param definition [Class] the name of a definition class
     # @param visit [Boolean] true if the context needs to be called into view
     # @param block [Proc] logic to execute within the context of the definition
     # @return [Object] instance of the definition
     def on(definition, visit=false, &block)
       definition = get_object_for(definition) if definition.is_a? String
       
-      return @active if @active.kind_of?(definition)
+      if @active.kind_of?(definition)
+        block.call @active if block
+        return @active
+      end
+      
       @active = definition.new(@driver, visit)
       block.call @active if block
       
       @active
     end
\ No newline at end of file