Refactoring puppet-test -- it now supports specifying a fork level (so you can get multiple tests running in parallel from one host), and tests are modeled explicitly, so it will be easier to add new tests.

git-svn-id: https://reductivelabs.com/svn/puppet/trunk@2623 980ebf18-57e1-0310-9a29-db15c13687c0

1 files changed, 139 insertions, 48 deletions

diff --git a/ext/puppet-test b/ext/puppet-test
index 35300b40b..8a82c4643 100755
--- a/ext/puppet-test
+++ b/ext/puppet-test
@@ -7,7 +7,7 @@
 # = Usage
 #
 #   puppet-test  [-c|--compile] [-D|--describe] [-d|--debug] [-f|--file <file>]
-#       [-h|--help] [-H|--hostname <host name>] [-r|--repeat <number=1>]
+#       [--fork <num>] [-h|--help] [-H|--hostname <host name>] [-r|--repeat <number=1>]
 #       [-V|--version] [-v|--verbose]
 #
 # = Description
@@ -24,7 +24,7 @@
 # See the configuration file documentation at
 # http://reductivelabs.com/projects/puppet/reference/configref.html for
 # the full list of acceptable parameters. A commented list of all
-# configuration options can also be generated by running puppetd with
+# configuration $options can also be generated by running puppetd with
 # '--genconfig'.
 #
 # compile::
@@ -40,6 +40,9 @@
 #   Test file performance.  Defaults to retrieving the file, but you can use
 #   --describe to instead describe the file.
 #
+# fork::
+#   Fork the specified number of times, thus acting as multiple clients.
+#
 # fqdn::
 #   Set the fully-qualified domain name of the client.  This is only used for
 #   certificate purposes, but can be used to override the discovered hostname.
@@ -80,10 +83,11 @@ require 'puppet'
 require 'puppet/network/client'
 require 'getoptlong'
 
-args = [
+$cmdargs = [
 	[ "--compile",    "-c",			GetoptLong::NO_ARGUMENT ],
 	[ "--describe",  "-D",			GetoptLong::NO_ARGUMENT ],
 	[ "--file",      "-f",			GetoptLong::REQUIRED_ARGUMENT ],
+	[ "--fork",          			GetoptLong::REQUIRED_ARGUMENT ],
 	[ "--fqdn",      "-F",			GetoptLong::REQUIRED_ARGUMENT ],
 	[ "--repeat",    "-r",			GetoptLong::REQUIRED_ARGUMENT ],
 	[ "--debug",	"-d",			GetoptLong::NO_ARGUMENT ],
@@ -92,15 +96,15 @@ args = [
 	[ "--version",	"-V",			GetoptLong::NO_ARGUMENT ],
 ]
 
-# Add all of the config parameters as valid options.
-Puppet.config.addargs(args)
+# Add all of the config parameters as valid $options.
+Puppet.config.addargs($cmdargs)
 Puppet::Util::Log.newdestination(:console)
 
-result = GetoptLong.new(*args)
+result = GetoptLong.new(*$cmdargs)
 
-args = {}
+$args = {}
 
-options = {:repeat => 1, :compile => true}
+$options = {:repeat => 1, :fork => 0}
 
 begin
     explicit_waitforcert = false
@@ -109,16 +113,24 @@ begin
             # First check to see if the argument is a valid configuration parameter;
             # if so, set it.
             when "--compile"
-                options[:compile] = true
+                $options[:test] = :compile
             when "--file"
-                options[:file] = arg
-                options[:compile] = false
+                $options[:file] = arg
+                # Allow --describe to set this.
+                $options[:test] ||= :retrieve
+            when "--fork"
+                begin
+                    $options[:fork] = Integer(arg)
+                rescue => detail
+                    $stderr.puts "The argument to 'fork' must be an integer"
+                    exit(14)
+                end
             when "--describe"
-                options[:describe] = true
+                $options[:test] = :describe
             when "--fqdn"
-                options[:fqdn] = arg
+                $options[:fqdn] = arg
             when "--repeat"
-                options[:repeat] = Integer(arg)
+                $options[:repeat] = Integer(arg)
             when "--help"
                 if Puppet.features.usage?
                     RDoc::usage && exit
@@ -148,54 +160,133 @@ end
 # Now parse the config
 Puppet.parse_config
 
-args[:Server] = Puppet[:server]
+$args[:Server] = Puppet[:server]
 
-if options[:compile]
-    args[:cache] = false
-    # Create a config client and pull the config down
-    client = Puppet::Network::Client.master.new(args)
-    unless client.read_cert
-        fail "Could not read client certificate"
-    end
+$options[:test] ||= :compile
 
-    # Use the facts from the cache, to skip the time it takes
-    # to load them.
-    client.dostorage
-    facts = Puppet::Util::Storage.cache(:configuration)[:facts]
+$pids = []
 
-    if facts.empty?
-        facts = client.master.getfacts
+class Test
+    attr_reader :name, :doc
+
+    def execute
+        raise "Test %s did not override 'execute'" % @name
     end
 
-    if host = options[:fqdn]
-        facts["fqdn"] = host
-        facts["hostname"] = host.sub(/\..+/, '')
-        facts["domain"] = host.sub(/^[^.]+\./, '')
+    def forked?
+        defined? @forking
     end
 
-    options[:repeat].times do
-        Puppet::Util.benchmark(:notice, "Retrieved configuration") do
-            client.send(:get_actual_config, facts)
-        end
+    def initialize(name, doc, &block)
+        @name = name
+        @doc = doc
+
+        raise "You must pass a block to the Test" unless block_given?
+        instance_eval(&block)
     end
-elsif options[:file]
-    # Create a config client and pull the config down
-    client = Puppet::Network::Client.file.new(args)
-    unless client.read_cert
-        fail "Could not read client certificate"
+
+    def prepare
+        raise "Test %s did not override 'prepare'" % @name
     end
 
-    options[:repeat].times do
-        if options[:describe]
-            Puppet::Util.benchmark(:notice, "Described file") do
-                client.describe(options[:file], :ignore)
+    # Run the actual test.
+    def run
+        puts "Running %s test" % @name
+        prepare()
+
+        if $options[:fork] > 0
+            @forking = true
+            $options[:fork].times {
+                if pid = fork
+                    $pids << pid
+                else
+                    break
+                end
+            }
+        end
+
+        $options[:repeat].times do
+            if forked?
+                msg = @doc + " in PID %s" % Process.pid
+            else
+                msg = @doc
             end
-        else
-            Puppet::Util.benchmark(:notice, "Retrieved file") do
-                client.retrieve(options[:file], :ignore)
+            Puppet::Util.benchmark(:notice, msg) do
+                execute()
             end
         end
     end
 end
 
+$tests = {}
+def newtest(name, desc,&block)
+    $tests[name] = Test.new(name, desc, &block)
+end
+
+newtest :compile, "Retrieved configuration" do
+    def prepare
+        $args[:cache] = false
+        # Create a config client and pull the config down
+        @client = Puppet::Network::Client.master.new($args)
+        unless @client.read_cert
+            fail "Could not read client certificate"
+        end
+
+        # Use the facts from the cache, to skip the time it takes
+        # to load them.
+        @client.dostorage
+        @facts = Puppet::Util::Storage.cache(:configuration)[:facts]
+
+        if @facts.empty?
+            @facts = @client.master.getfacts
+        end
+
+        if host = $options[:fqdn]
+            @facts["fqdn"] = host
+            @facts["hostname"] = host.sub(/\..+/, '')
+            @facts["domain"] = host.sub(/^[^.]+\./, '')
+        end
+    end
+
+    def execute
+        @client.send(:get_actual_config, @facts)
+    end
+end
+
+newtest :describe, "Described file" do
+    def prepare
+        @client = Puppet::Network::Client.file.new($args)
+        unless @client.read_cert
+            fail "Could not read client certificate"
+        end
+    end
+
+    def execute
+        @client.describe($options[:file], :ignore)
+    end
+end
+
+newtest :retrieve, "Retrieved file" do
+    def prepare
+        @client = Puppet::Network::Client.file.new($args)
+        unless @client.read_cert
+            fail "Could not read client certificate"
+        end
+    end
+
+    def execute
+        client.retrieve($options[:file], :ignore)
+    end
+end
+
+unless test = $tests[$options[:test]]
+    raise "Invalid test %s; valid tests are %s" % [options[:test], $tests.keys.collect { |t| t.to_s }.join(", ")]
+end
+
+test.run
+
+if $options[:fork] > 0
+    Process.waitall
+end
+
 # $Id$