Adding puppet-test, which is useful for performance testing

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

1 files changed, 201 insertions, 0 deletions

diff --git a/ext/puppet-test b/ext/puppet-test
new file mode 100755
index 000000000..35300b40b
--- /dev/null
+++ b/ext/puppet-test
@@ -0,0 +1,201 @@
+#!/usr/bin/env ruby
+# == Synopsis 
+#
+# Test individual client performance.  Can compile configurations, describe
+# files, or retrieve files.
+#
+# = Usage
+#
+#   puppet-test  [-c|--compile] [-D|--describe] [-d|--debug] [-f|--file <file>]
+#       [-h|--help] [-H|--hostname <host name>] [-r|--repeat <number=1>]
+#       [-V|--version] [-v|--verbose]
+#
+# = Description
+#
+# This is a simple script meant for doing performance tests with Puppet.  By
+# default it pulls down a compiled configuration.
+#
+# = Options
+#
+# Note that any configuration parameter that's valid in the configuration file
+# is also a valid long argument.  For example, 'server' is a valid configuration
+# parameter, so you can specify '--server <servername>' as an argument.
+#
+# 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
+# '--genconfig'.
+#
+# compile::
+#   Compile the client's configuration.  The default.
+#
+# debug::
+#   Enable full debugging.
+#
+# describe::
+#   Describe the file being tested, rather than retrieve it.
+#
+# file::
+#   Test file performance.  Defaults to retrieving the file, but you can use
+#   --describe to instead describe the file.
+#
+# 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.
+#   If you need to use this flag, it is generally an indication of a setup problem.
+#
+# help::
+#   Print this help message
+#
+# repeat::
+#  How many times to perform the test.
+#
+# verbose::
+#   Turn on verbose reporting.
+#
+# version::
+#   Print the puppet version number and exit.
+#
+# = Example
+#
+#   puppet-test --file /module/path/to/file
+#
+# = Author
+#
+# Luke Kanies
+#
+# = Copyright
+#
+# Copyright (c) 2005, 2006 Reductive Labs, LLC
+# Licensed under the GNU Public License
+
+# Do an initial trap, so that cancels don't get a stack trace.
+trap(:INT) do
+    $stderr.puts "Cancelling startup"
+    exit(0)
+end
+
+require 'puppet'
+require 'puppet/network/client'
+require 'getoptlong'
+
+args = [
+	[ "--compile",    "-c",			GetoptLong::NO_ARGUMENT ],
+	[ "--describe",  "-D",			GetoptLong::NO_ARGUMENT ],
+	[ "--file",      "-f",			GetoptLong::REQUIRED_ARGUMENT ],
+	[ "--fqdn",      "-F",			GetoptLong::REQUIRED_ARGUMENT ],
+	[ "--repeat",    "-r",			GetoptLong::REQUIRED_ARGUMENT ],
+	[ "--debug",	"-d",			GetoptLong::NO_ARGUMENT ],
+	[ "--help",		"-h",			GetoptLong::NO_ARGUMENT ],
+	[ "--verbose",	"-v",			GetoptLong::NO_ARGUMENT ],
+	[ "--version",	"-V",			GetoptLong::NO_ARGUMENT ],
+]
+
+# Add all of the config parameters as valid options.
+Puppet.config.addargs(args)
+Puppet::Util::Log.newdestination(:console)
+
+result = GetoptLong.new(*args)
+
+args = {}
+
+options = {:repeat => 1, :compile => true}
+
+begin
+    explicit_waitforcert = false
+    result.each { |opt,arg|
+        case opt
+            # First check to see if the argument is a valid configuration parameter;
+            # if so, set it.
+            when "--compile"
+                options[:compile] = true
+            when "--file"
+                options[:file] = arg
+                options[:compile] = false
+            when "--describe"
+                options[:describe] = true
+            when "--fqdn"
+                options[:fqdn] = arg
+            when "--repeat"
+                options[:repeat] = Integer(arg)
+            when "--help"
+                if Puppet.features.usage?
+                    RDoc::usage && exit
+                else
+                    puts "No help available unless you have RDoc::usage installed"
+                    exit
+                end
+            when "--version"
+                puts "%s" % Puppet.version
+                exit
+            when "--verbose"
+                Puppet::Util::Log.level = :info
+                Puppet::Util::Log.newdestination(:console)
+            when "--debug"
+                Puppet::Util::Log.level = :debug
+                Puppet::Util::Log.newdestination(:console)
+            else
+                Puppet.config.handlearg(opt, arg)
+        end
+    }
+rescue GetoptLong::InvalidOption => detail
+    $stderr.puts detail
+    $stderr.puts "Try '#{$0} --help'"
+    exit(1)
+end
+
+# Now parse the config
+Puppet.parse_config
+
+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
+
+    # 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
+
+    options[:repeat].times do
+        Puppet::Util.benchmark(:notice, "Retrieved configuration") do
+            client.send(:get_actual_config, facts)
+        end
+    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"
+    end
+
+    options[:repeat].times do
+        if options[:describe]
+            Puppet::Util.benchmark(:notice, "Described file") do
+                client.describe(options[:file], :ignore)
+            end
+        else
+            Puppet::Util.benchmark(:notice, "Retrieved file") do
+                client.retrieve(options[:file], :ignore)
+            end
+        end
+    end
+end
+
+# $Id$