lib/floe/cli.rb in floe-0.15.0 vs lib/floe/cli.rb in floe-0.15.1

@@ -1,11 +1,14 @@
+require "floe"
+require "floe/container_runner"
+
 module Floe
   class CLI
+    include Logging
+
     def initialize
       require "optimist"
-      require "floe"
-      require "floe/container_runner"
       require "logger"
 
       Floe.logger = Logger.new($stdout)
       Floe.logger.level = 0 if ENV["DEBUG"]
     end
@@ -18,16 +21,26 @@
       workflows =
         workflows_inputs.each_slice(2).map do |workflow, input|
           create_workflow(workflow, opts[:context], input, credentials)
         end
 
-      Floe::Workflow.wait(workflows, &:run_nonblock)
+      output_streams = create_loggers(workflows, opts[:segment_output])
 
+      logger.info("Checking #{workflows.count} workflows...")
+      ready = Floe::Workflow.wait(workflows, &:run_nonblock)
+      logger.info("Checking #{workflows.count} workflows...Complete - #{ready.count} ready")
+
       # Display status
       workflows.each do |workflow|
-        puts "", "#{workflow.name}#{" (#{workflow.status})" unless workflow.context.success?}", "===" if workflows.size > 1
-        puts workflow.output
+        if workflows.size > 1
+          logger.info("")
+          logger.info("#{workflow.name}#{" (#{workflow.status})" unless workflow.context.success?}")
+          logger.info("===")
+        end
+
+        logger.info(output_streams[workflow].string) if output_streams[workflow]
+        logger.info(workflow.output)
       end
 
       workflows.all? { |workflow| workflow.context.success? }
     rescue Floe::Error => err
       abort(err.message)
@@ -47,10 +60,11 @@
             If --input is not passed and --workflow is passed, defaults to '{}'.
         EOMSG
         opt :context, "JSON payload of the Context",              :type => :string
         opt :credentials, "JSON payload with Credentials",        :type => :string
         opt :credentials_file, "Path to a file with Credentials", :type => :string
+        opt :segment_output, "Segment output by each worker",     :default => false
 
         Floe::ContainerRunner.cli_options(self)
 
         banner("")
         banner("General options:")
@@ -86,8 +100,20 @@
     end
 
     def create_workflow(workflow, context_payload, input, credentials)
       context = Floe::Workflow::Context.new(context_payload, :input => input, :credentials => credentials)
       Floe::Workflow.load(workflow, context)
+    end
+
+    def create_loggers(workflows, segment_output)
+      if workflows.size == 1 || !segment_output
+        # no extra work necessary
+        {}
+      else
+        workflows.each_with_object({}) do |workflow, h|
+          workflow.context.logger = Logger.new(output = StringIO.new)
+          h[workflow] = output
+        end
+      end
     end
   end
 end