Move puppetrun to Application Controller paradigm

Signed-off-by: Brice Figureau <brice-puppet@daysofwonder.com>

3 files changed, 494 insertions, 241 deletions

diff --git a/bin/puppetrun b/bin/puppetrun
index 28f72d9c3..da4f24a98 100755
--- a/bin/puppetrun
+++ b/bin/puppetrun
@@ -125,245 +125,6 @@
 # Copyright (c) 2005 Reductive Labs, LLC
 # Licensed under the GNU Public License
 
-[:INT, :TERM].each do |signal|
-    trap(signal) do
-        $stderr.puts "Cancelling"
-        exit(1)
-    end
-end
-
-begin
-    require 'rubygems'
-rescue LoadError
-    # Nothing; we were just doing this just in case
-end
-
-begin
-    require 'ldap'
-rescue LoadError
-    $stderr.puts "Failed to load ruby LDAP library. LDAP functionality will not be available"
-end
-
-require 'puppet'
-require 'puppet/network/client'
-require 'puppet/util/ldap/connection'
-require 'getoptlong'
-
-flags = [
-    [ "--all",      "-a",	        GetoptLong::NO_ARGUMENT ],
-    [ "--tag",      "-t",           GetoptLong::REQUIRED_ARGUMENT ],
-    [ "--class",	"-c",			GetoptLong::REQUIRED_ARGUMENT ],
-    [ "--foreground", "-f",	        GetoptLong::NO_ARGUMENT ],
-    [ "--debug",	"-d",			GetoptLong::NO_ARGUMENT ],
-    [ "--help",		"-h",			GetoptLong::NO_ARGUMENT ],
-    [ "--host",         			GetoptLong::REQUIRED_ARGUMENT ],
-    [ "--parallel", "-p",			GetoptLong::REQUIRED_ARGUMENT ],
-    [ "--ping", "-P",			GetoptLong::NO_ARGUMENT ],
-    [ "--no-fqdn",  "-n",			GetoptLong::NO_ARGUMENT ],
-    [ "--test",                     GetoptLong::NO_ARGUMENT ],
-    [ "--version",  "-V",           GetoptLong::NO_ARGUMENT ]
-]
-
-# Add all of the config parameters as valid options.
-Puppet.settings.addargs(flags)
-
-result = GetoptLong.new(*flags)
-
-options = {
-    :ignoreschedules => false,
-    :foreground => false,
-    :parallel => 1,
-    :ping => false,
-    :debug => false,
-    :test => false,
-    :all => false,
-    :verbose => true,
-    :fqdn => true
-}
-
-hosts = []
-classes = []
-tags = []
-
-Puppet::Util::Log.newdestination(:console)
-
-begin
-    result.each { |opt,arg|
-        case opt
-            when "--version"
-                puts "%s" % Puppet.version
-                exit
-            when "--ignoreschedules"
-                options[:ignoreschedules] = true
-            when "--no-fqdn"
-                options[:fqdn] = false
-            when "--all"
-                options[:all] = true
-            when "--test"
-                options[:test] = true
-            when "--tag"
-                tags << arg
-            when "--class"
-                classes << arg
-            when "--host"
-                hosts << arg
-            when "--help"
-                if Puppet.features.usage?
-                    RDoc::usage && exit
-                else
-                    puts "No help available unless you have RDoc::usage installed"
-                    exit
-                end
-            when "--parallel"
-                begin
-                    options[:parallel] = Integer(arg)
-                rescue
-                    $stderr.puts "Could not convert %s to an integer" % arg.inspect
-                    exit(23)
-                end
-            when "--ping"
-                 options[:ping] = true 
-            when "--foreground"
-                options[:foreground] = true
-            when "--debug"
-                options[:debug] = true
-            else
-                Puppet.settings.handlearg(opt, arg)
-        end
-    }
-rescue GetoptLong::InvalidOption => detail
-    $stderr.puts "Try '#{$0} --help'"
-    exit(1)
-end
-
-if options[:debug]
-    Puppet::Util::Log.level = :debug
-else
-    Puppet::Util::Log.level = :info
-end
-
-# Now parse the config
-Puppet.parse_config
-
-if Puppet[:node_terminus] == "ldap" and (options[:all] or classes)
-    if options[:all]
-        hosts = Puppet::Node.search("whatever").collect { |node| node.name }
-        puts "all: %s" % hosts.join(", ")
-    else
-        hosts = []
-        classes.each do |klass|
-            list = Puppet::Node.search("whatever", :class => klass).collect { |node| node.name }
-            puts "%s: %s" % [klass, list.join(", ")]
-
-            hosts += list
-        end
-    end
-elsif ! classes.empty?
-    $stderr.puts "You must be using LDAP to specify host classes"
-    exit(24)
-end
-
-if tags.empty?
-    tags = ""
-else
-    tags = tags.join(",")
-end
-
-children = {}
-
-# If we get a signal, then kill all of our children and get out.
-[:INT, :TERM].each do |signal|
-    trap(signal) do
-        Puppet.notice "Caught #{signal}; shutting down"
-        children.each do |pid, host|
-            Process.kill("INT", pid)
-        end
-
-        waitall
-
-        exit(1)
-    end
-end
-
-if options[:test]
-    puts "Skipping execution in test mode"
-    exit(0)
-end
-
-todo = hosts.dup
-
-failures = []
-
-# Now do the actual work
-go = true
-while go
-    # If we don't have enough children in process and we still have hosts left to
-    # do, then do the next host.
-    if children.length < options[:parallel] and ! todo.empty?
-        host = todo.shift
-        pid = fork do
-            if options[:ping]
-              out = %x{ping -c 1 #{host}}
-              unless $? == 0
-                  $stderr.print "Could not contact %s\n" % host
-                  next
-              end
-            end
-            client = Puppet::Network::Client.runner.new(
-                :Server => host,
-                :Port => Puppet[:puppetport]
-            )
-
-            print "Triggering %s\n" % host
-            begin
-                result = client.run(tags, options[:ignoreschedules], options[:foreground])
-            rescue => detail
-                $stderr.puts "Host %s failed: %s\n" % [host, detail]
-                exit(2)
-            end
-            
-            case result
-            when "success": exit(0)
-            when "running":
-                $stderr.puts "Host %s is already running" % host
-                exit(3)
-            else
-                $stderr.puts "Host %s returned unknown answer '%s'" % [host, result]
-                exit(12)
-            end
-        end
-        children[pid] = host
-    else
-        # Else, see if we can reap a process.
-        begin
-            pid = Process.wait
-
-            if host = children[pid]
-                # Remove our host from the list of children, so the parallelization
-                # continues working.
-                children.delete(pid)
-                if $?.exitstatus != 0
-                    failures << host
-                end
-                print "%s finished with exit code %s\n" % [host, $?.exitstatus]
-            else
-                $stderr.puts "Could not find host for PID %s with status %s" %
-                    [pid, $?.exitstatus]
-            end
-        rescue Errno::ECHILD
-            # There are no children left, so just exit unless there are still
-            # children left to do.
-            next unless todo.empty?
-
-            if failures.empty?
-                puts "Finished"
-                exit(0)
-            else
-                puts "Failed: %s" % failures.join(", ")
-                exit(3)
-            end
-        end
-    end
-end
-
+require 'puppet/application/puppetrun'
+Puppet::Application[:puppetrun].run
 
diff --git a/lib/puppet/application/puppetrun.rb b/lib/puppet/application/puppetrun.rb
new file mode 100644
index 000000000..9e9d2b81f
--- /dev/null
+++ b/lib/puppet/application/puppetrun.rb
@@ -0,0 +1,221 @@
+begin
+    require 'rubygems'
+rescue LoadError
+    # Nothing; we were just doing this just in case
+end
+
+begin
+    require 'ldap'
+rescue LoadError
+    $stderr.puts "Failed to load ruby LDAP library. LDAP functionality will not be available"
+end
+
+require 'puppet'
+require 'puppet/application'
+
+Puppet::Application.new(:puppetrun) do
+
+    should_not_parse_config
+
+    attr_accessor :hosts, :tags, :classes
+
+    option("--all","-a")
+    option("--foreground","-f")
+    option("--debug","-d")
+    option("--ping","-P")
+    option("--test")
+
+    option("--version", "-V") do |arg|
+        puts "%s" % Puppet.version
+        exit
+    end
+
+    option("--host HOST") do |arg|
+        @hosts << arg
+    end
+
+    option("--tag TAG", "-t") do |arg|
+        @tags << arg
+    end
+
+    option("--class CLASS", "-c") do |arg|
+        @classes << arg
+    end
+
+    option("--no-fqdn", "-n") do |arg|
+        options[:fqdn] = false
+    end
+
+    option("--parallel PARALLEL", "-p") do |arg|
+        begin
+            options[:parallel] = Integer(arg)
+        rescue
+            $stderr.puts "Could not convert %s to an integer" % arg.inspect
+            exit(23)
+        end
+    end
+
+
+    dispatch do
+        options[:test] ? :test : :main
+    end
+
+    command(:test) do
+        puts "Skipping execution in test mode"
+        exit(0)
+    end
+
+    command(:main) do
+        require 'puppet/network/client'
+        require 'puppet/util/ldap/connection'
+
+        todo = @hosts.dup
+
+        failures = []
+
+        # Now do the actual work
+        go = true
+        while go
+            # If we don't have enough children in process and we still have hosts left to
+            # do, then do the next host.
+            if @children.length < options[:parallel] and ! todo.empty?
+                host = todo.shift
+                pid = fork do
+                    run_for_host(host)
+                end
+                @children[pid] = host
+            else
+                # Else, see if we can reap a process.
+                begin
+                    pid = Process.wait
+
+                    if host = @children[pid]
+                        # Remove our host from the list of children, so the parallelization
+                        # continues working.
+                        @children.delete(pid)
+                        if $?.exitstatus != 0
+                            failures << host
+                        end
+                        print "%s finished with exit code %s\n" % [host, $?.exitstatus]
+                    else
+                        $stderr.puts "Could not find host for PID %s with status %s" %
+                            [pid, $?.exitstatus]
+                    end
+                rescue Errno::ECHILD
+                    # There are no children left, so just exit unless there are still
+                    # children left to do.
+                    next unless todo.empty?
+
+                    if failures.empty?
+                        puts "Finished"
+                        exit(0)
+                    else
+                        puts "Failed: %s" % failures.join(", ")
+                        exit(3)
+                    end
+                end
+            end
+        end
+    end
+
+    def run_for_host(host)
+        if options[:ping]
+            out = %x{ping -c 1 #{host}}
+            unless $? == 0
+                $stderr.print "Could not contact %s\n" % host
+                next
+            end
+        end
+        client = Puppet::Network::Client.runner.new(
+            :Server => host,
+            :Port => Puppet[:puppetport]
+        )
+
+        print "Triggering %s\n" % host
+        begin
+            result = client.run(@tags, options[:ignoreschedules], options[:foreground])
+        rescue => detail
+            $stderr.puts "Host %s failed: %s\n" % [host, detail]
+            exit(2)
+        end
+
+        case result
+        when "success": exit(0)
+        when "running":
+            $stderr.puts "Host %s is already running" % host
+            exit(3)
+        else
+            $stderr.puts "Host %s returned unknown answer '%s'" % [host, result]
+            exit(12)
+        end
+    end
+
+    preinit do
+        [:INT, :TERM].each do |signal|
+            trap(signal) do
+                $stderr.puts "Cancelling"
+                exit(1)
+            end
+        end
+        options[:parallel] = 1
+        options[:verbose] = true
+        options[:fqdn] = true
+
+        @hosts = []
+        @classes = []
+        @tags = []
+    end
+
+    setup do
+        if options[:debug]
+            Puppet::Util::Log.level = :debug
+        else
+            Puppet::Util::Log.level = :info
+        end
+
+        # Now parse the config
+        Puppet.parse_config
+
+        if Puppet[:node_terminus] == "ldap" and (options[:all] or @classes)
+            if options[:all]
+                @hosts = Puppet::Node.search("whatever").collect { |node| node.name }
+                puts "all: %s" % @hosts.join(", ")
+            else
+                @hosts = []
+                @classes.each do |klass|
+                    list = Puppet::Node.search("whatever", :class => klass).collect { |node| node.name }
+                    puts "%s: %s" % [klass, list.join(", ")]
+
+                    @hosts += list
+                end
+            end
+        elsif ! @classes.empty?
+            $stderr.puts "You must be using LDAP to specify host classes"
+            exit(24)
+        end
+
+        if @tags.empty?
+            @tags = ""
+        else
+            @tags = @tags.join(",")
+        end
+
+        @children = {}
+
+        # If we get a signal, then kill all of our children and get out.
+        [:INT, :TERM].each do |signal|
+            trap(signal) do
+                Puppet.notice "Caught #{signal}; shutting down"
+                @children.each do |pid, host|
+                    Process.kill("INT", pid)
+                end
+
+                waitall
+
+                exit(1)
+            end
+        end
+
+    end
+
+end
diff --git a/spec/unit/application/puppetrun.rb b/spec/unit/application/puppetrun.rb
new file mode 100755
index 000000000..5d033e960
--- /dev/null
+++ b/spec/unit/application/puppetrun.rb
@@ -0,0 +1,271 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../../spec_helper'
+
+require 'puppet/util/ldap/connection'
+require 'puppet/application/puppetrun'
+
+describe "puppetrun" do
+    before :each do
+        Puppet::Util::Ldap::Connection.stubs(:new).returns(stub_everything)
+        @puppetrun = Puppet::Application[:puppetrun]
+    end
+
+    it "should ask Puppet::Application to not parse Puppet configuration file" do
+        @puppetrun.should_parse_config?.should be_false
+    end
+
+    it "should declare a main command" do
+        @puppetrun.should respond_to(:main)
+    end
+
+    it "should declare a test command" do
+        @puppetrun.should respond_to(:test)
+    end
+
+    it "should declare a preinit block" do
+        @puppetrun.should respond_to(:run_preinit)
+    end
+
+    describe "during preinit" do
+        before :each do
+            @puppetrun.stubs(:trap)
+        end
+
+        it "should catch INT and TERM" do
+            @puppetrun.stubs(:trap).with { |arg,block| arg == :INT or arg == :TERM }
+
+            @puppetrun.run_preinit
+        end
+
+        it "should set parallel option to 1" do
+            @puppetrun.run_preinit
+
+            @puppetrun.options[:parallel].should == 1
+        end
+
+        it "should set verbose by default" do
+            @puppetrun.run_preinit
+
+            @puppetrun.options[:verbose].should be_true
+        end
+
+        it "should set fqdn by default" do
+            @puppetrun.run_preinit
+
+            @puppetrun.options[:fqdn].should be_true
+        end
+    end
+
+    describe "when applying options" do
+
+        [:all, :foreground, :debug, :ping, :test].each do |option|
+            it "should declare handle_#{option} method" do
+                @puppetrun.should respond_to("handle_#{option}".to_sym)
+            end
+
+            it "should store argument value when calling handle_#{option}" do
+                @puppetrun.options.expects(:[]=).with(option, 'arg')
+                @puppetrun.send("handle_#{option}".to_sym, 'arg')
+            end
+        end
+
+        it "should exit after printing the version" do
+            @puppetrun.stubs(:puts)
+
+            lambda { @puppetrun.handle_version(nil) }.should raise_error(SystemExit)
+        end
+
+        it "should add to the host list with the host option" do
+            @puppetrun.handle_host('host')
+
+            @puppetrun.hosts.should == ['host']
+        end
+
+        it "should add to the tag list with the tag option" do
+            @puppetrun.handle_tag('tag')
+
+            @puppetrun.tags.should == ['tag']
+        end
+
+        it "should add to the class list with the class option" do
+            @puppetrun.handle_class('class')
+
+            @puppetrun.classes.should == ['class']
+        end
+    end
+
+    describe "during setup" do
+
+        before :each do
+            @puppetrun.classes = []
+            @puppetrun.tags = []
+            @puppetrun.hosts = []
+            Puppet::Log.stubs(:level=)
+            @puppetrun.stubs(:trap)
+            @puppetrun.stubs(:puts)
+            Puppet.stubs(:parse_config)
+
+            @puppetrun.options.stubs(:[]).with(any_parameters)
+        end
+
+        it "should set log level to debug if --debug was passed" do
+            @puppetrun.options.stubs(:[]).with(:debug).returns(true)
+
+            Puppet::Log.expects(:level=).with(:debug)
+
+            @puppetrun.run_setup
+        end
+
+        it "should set log level to info if --verbose was passed" do
+            @puppetrun.options.stubs(:[]).with(:verbose).returns(true)
+
+            Puppet::Log.expects(:level=).with(:info)
+
+            @puppetrun.run_setup
+        end
+
+        it "should Parse puppet config" do
+            Puppet.expects(:parse_config)
+
+            @puppetrun.run_setup
+        end
+
+        describe "when using the ldap node terminus" do
+            before :each do
+                Puppet.stubs(:[]).with(:node_terminus).returns("ldap")
+            end
+
+            it "should search for all nodes if --all" do
+                @puppetrun.options.stubs(:[]).with(:all).returns(true)
+                @puppetrun.stubs(:puts)
+
+                Puppet::Node.expects(:search).with("whatever").returns([])
+
+                @puppetrun.run_setup
+            end
+
+            it "should search for nodes including given classes" do
+                @puppetrun.options.stubs(:[]).with(:all).returns(false)
+                @puppetrun.stubs(:puts)
+                @puppetrun.classes = ['class']
+
+                Puppet::Node.expects(:search).with("whatever", :class => "class").returns([])
+
+                @puppetrun.run_setup
+            end
+        end
+
+        describe "when using regular nodes" do
+            it "should fail if some classes have been specified" do
+                $stderr.stubs(:puts)
+                @puppetrun.classes = ['class']
+
+                @puppetrun.expects(:exit).with(24)
+
+                @puppetrun.run_setup
+            end
+        end
+    end
+
+    describe "when running" do
+        before :each do
+            @puppetrun.stubs(:puts)
+        end
+
+        it "should dispatch to test if --test is used" do
+            @puppetrun.options.stubs(:[]).with(:test).returns(true)
+
+            @puppetrun.get_command.should == :test
+        end
+
+        it "should dispatch to main if --test is not used" do
+            @puppetrun.options.stubs(:[]).with(:test).returns(false)
+
+            @puppetrun.get_command.should == :main
+        end
+
+        describe "the test command" do
+            it "should exit with exit code 0 " do
+                @puppetrun.expects(:exit).with(0)
+
+                @puppetrun.test
+            end
+        end
+
+        describe "the main command" do
+            before :each do
+                @client = stub_everything 'client'
+                @client.stubs(:run).returns("success")
+                Puppet::Network::Client.runner.stubs(:new).returns(@client)
+                @puppetrun.options.stubs(:[]).with(:parallel).returns(1)
+                @puppetrun.options.stubs(:[]).with(:ping).returns(false)
+                @puppetrun.options.stubs(:[]).with(:ignoreschedules).returns(false)
+                @puppetrun.options.stubs(:[]).with(:foreground).returns(false)
+                @puppetrun.stubs(:print)
+                @puppetrun.stubs(:exit)
+                $stderr.stubs(:puts)
+            end
+
+            it "should create as much childs as --parallel" do
+                @puppetrun.options.stubs(:[]).with(:parallel).returns(3)
+                @puppetrun.hosts = ['host1', 'host2', 'host3']
+                @puppetrun.stubs(:exit).raises(SystemExit)
+                Process.stubs(:wait).returns(1).then.returns(2).then.returns(3).then.raises(Errno::ECHILD)
+
+                @puppetrun.expects(:fork).times(3).returns(1).then.returns(2).then.returns(3)
+
+                lambda { @puppetrun.main }.should raise_error
+            end
+
+            it "should delegate to run_for_host per host" do
+                @puppetrun.hosts = ['host1', 'host2']
+                @puppetrun.stubs(:exit).raises(SystemExit)
+                @puppetrun.stubs(:fork).returns(1).yields
+                Process.stubs(:wait).returns(1).then.raises(Errno::ECHILD)
+
+                @puppetrun.expects(:run_for_host).times(2)
+
+                lambda { @puppetrun.main }.should raise_error
+            end
+
+            describe "during call of run_for_host" do
+                it "should create a Runner Client per given host" do
+                    Puppet::Network::Client.runner.expects(:new).returns(@client)
+
+                    @puppetrun.run_for_host('host')
+                end
+
+                it "should call Client.run for the given host" do
+                    @client.expects(:run)
+
+                    @puppetrun.run_for_host('host')
+                end
+
+                it "should exit the child with 0 on success" do
+                    @client.stubs(:run).returns("success")
+
+                    @puppetrun.expects(:exit).with(0)
+
+                    @puppetrun.run_for_host('host')
+                end
+
+                it "should exit the child with 3 on running" do
+                    @client.stubs(:run).returns("running")
+
+                    @puppetrun.expects(:exit).with(3)
+
+                    @puppetrun.run_for_host('host')
+                end
+
+                it "should exit the child with 12 on unknown answer" do
+                    @client.stubs(:run).returns("whatever")
+
+                    @puppetrun.expects(:exit).with(12)
+
+                    @puppetrun.run_for_host('host')
+                end
+            end
+        end
+    end
+end