#!/usr/bin/env ruby
# == Synopsis
#
# Test individual client performance.  Can compile configurations, describe
# files, or retrieve files.
#
# = Usage
#
#   puppet-test  [-c|--compile] [-D|--describe <file>] [-d|--debug]
#       [--fork <num>] [-h|--help] [-H|--hostname <host name>] [-l|--list] [-r|--repeat <number=1>]
#       [-R|--retrieve <file>] [-t|--test <test>] [-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, but it supports multiple
# other tests.
#
# = 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.  This is a query to get information about
#   the file from the server, to determine if it should be copied, and is the
#   first half of every file transfer.
#
# 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.
#   If you need to use this flag, it is generally an indication of a setup problem.
#
# help::
#   Print this help message
#
# list::
#   List all available tests.
#
# node::
#  Specify the node to use.  This is useful for looking up cached yaml data
#  in your :clientyaml directory, and forcing a specific host's configuration to
#  get compiled.
#
# pause::
#  Pause before starting test (useful for testing with dtrace).
#
# repeat::
#  How many times to perform the test.
#
# retrieve::
#   Test file retrieval performance.  Retrieves the specified file from the
#   remote system.  Note that the server should be specified via --server,
#   so the argument to this option is just the remote module name and path,
#   e.g., "/dist/apps/myapp/myfile", where "dist" is the module and
#   "apps/myapp/myfile" is the path to the file relative to the module.
#
# test::
#   Specify the test to run.  You can see the list of tests by running this command with --list.
#
# verbose::
#   Turn on verbose reporting.
#
# version::
#   Print the puppet version number and exit.
#
# = Example
#
#   puppet-test --retrieve /module/path/to/file
#
# = License
#   Copyright 2011 Luke Kanies
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

# Do an initial trap, so that cancels don't get a stack trace.
trap(:INT) do
    $stderr.puts "Cancelling startup"
    exit(1)
end

require 'puppet'
require 'puppet/network/client'
require 'getoptlong'

class Suite
    attr_reader :name, :doc

    @@suites = {}
    @@tests = {}

    def self.[](name)
        @@suites[name]
    end

    # Run a test by first finding the suite then running the appropriate test.
    def self.run(test)
        unless suite_name = @@tests[test]
            raise "Unknown test %s" % test

        end
        unless suite = @@suites[suite_name]
            raise "Unknown test suite %s from test %s" % [suite_name, test]
        end

        suite.run(test)
    end

    # What suites are available?
    def self.suites
        @@suites.keys
    end

    def forked?
        defined? @forking
    end

    # Create a new test suite.
    def initialize(name, doc, &block)
        @name = name
        @doc = doc

        @tests = {}

        @@suites[name] = self

        raise "You must pass a block to the Test" unless block_given?
        instance_eval(&block)
    end

    # Define a new type of test on this suite.
    def newtest(name, doc, &block)
        @tests[name] = doc

        if @@tests[name]
            raise "Test names must be unique; cannot redefine %s" % name
        end

        @@tests[name] = @name

        meta_def(name, &block)
    end

    # Run the actual test.
    def run(test)
        unless doc = @tests[test]
            raise "Suite %s only supports tests %s; not %s" % [@name, @tests.keys.collect { |k| k.to_s }.join(","), test]
        end
        puts "Running %s %s test" % [@name, test]

        prepare() if respond_to?(:prepare)

        if $options[:pause]
            puts "Hit any key to continue"
            $stdin.readline
            puts "Continuing with test"
        end

        if $options[:fork] > 0
            @forking = true
            $options[:fork].times {
                if pid = fork
                    $pids << pid
                else
                    break
                end
            }
        end

        $options[:repeat].times do |i|
            @count = i
            if forked?
                msg = doc + " in PID %s" % Process.pid
            else
                msg = doc
            end
            Puppet::Util.benchmark(:notice, msg) do
                begin
                    send(test)
                rescue => detail
                    puts detail.backtrace if Puppet[:trace]
                    Puppet.err "%s failed: %s" % [@name, detail.to_s]
                end
            end
        end
    end

    # What tests are available on this suite?
    def tests
        @tests.keys
    end
end

Suite.new :parser, "Manifest parsing" do
    newtest :parse, "Parsed files" do
        @parser = Puppet::Parser::Parser.new(:environment => Puppet[:environment])
        @parser.file = Puppet[:manifest]
        @parser.parse
    end
end

Suite.new :local_catalog, "Local catalog handling" do
    def prepare
        @node = Puppet::Node.find($options[:nodes][0])
    end

    newtest :compile, "Compiled catalog" do
        Puppet::Resource::Catalog.find(@node)
    end
end

Suite.new :resource_type, "Managing resource types" do
    newtest :find, "Find a type" do
        Puppet::Resource::Type.terminus_class = :parser
        ARGV.each do |name|
            json = Puppet::Resource::Type.find(name).to_pson
            data = PSON.parse(json)
            p Puppet::Resource::Type.from_pson(data)
        end
    end

    newtest :search_types, "Find all types" do
        Puppet::Resource::Type.terminus_class = :rest
        result = Puppet::Resource::Type.search("*")
        result.each { |r| p r }
    end

    newtest :restful_type, "Find a type and return it via REST" do
        Puppet::Resource::Type.terminus_class = :rest
        ARGV.each do |name|
            p Puppet::Resource::Type.find(name)
        end
    end
end

Suite.new :remote_catalog, "Remote catalog handling" 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

        if tmp = Puppet::Node::Facts.find($options[:nodes][0])
            @facts = tmp.values
        else
            raise "Could not find facts for %s" % $optins[:nodes][0]
        end

        if host = $options[:fqdn]
            @facts["fqdn"] = host
            @facts["hostname"] = host.sub(/\..+/, '')
            @facts["domain"] = host.sub(/^[^.]+\./, '')
        end

        @facts = YAML.dump(@facts)
    end

    newtest :getconfig, "Compiled catalog" do
        @client.driver.getconfig(@facts, "yaml")
    end

    # This test will always force a false answer.
    newtest :fresh, "Checked freshness" do
        @client.driver.freshness
    end
end

Suite.new :file, "File interactions" do
    def prepare
        unless $options[:file]
            fail "You must specify a file (using --file <file>) to interact with on the server"
        end
        @client = Puppet::Network::Client.file.new($args)
        unless @client.read_cert
            fail "Could not read client certificate"
        end
    end

    newtest :describe, "Described file" do
        @client.describe($options[:file], :ignore)
    end

    newtest :retrieve, "Retrieved file" do
        @client.retrieve($options[:file], :ignore)
    end
end

Suite.new :filebucket, "Filebucket interactions" do
    def prepare
        require 'tempfile'
        @client = Puppet::FileBucket::Dipper.new($args)
    end

    newtest :backup, "Backed up file" do
        Tempfile.open("bucket_testing") do |f|
            f.print rand(1024)
            f.close
            @client.backup(f.path)
        end
    end
end

# Note that this uses an env variable to determine how many resources per
# host to create (with a default of 10).  'repeat' determines how
# many hosts to create.  You probably will get mad failures if you
# use forking and sqlite.
# Here's an example run of this test, using sqlite:
#   RESOURCE_COUNT=50 ext/puppet-test --suite rails --test storage --confdir /tmp/storagetesting --vardir /tmp/storagetesting --repeat 10
Suite.new :rails, "Rails Interactions" do
    def prepare
        Puppet::Rails.init
        @facts = Facter.to_hash

        if num = ENV["RESOURCECOUNT"]
            @resources = Integer(num)
        else
            @resources = 10
        end
    end

    Resource = Puppet::Parser::Resource

    def execute(test, msg)
        begin
            send(test)
        rescue => detail
            puts detail.backtrace if Puppet[:trace]
            Puppet.err "%s failed: %s" % [@name, detail.to_s]
        end
    end

    def mkresource(type, title, parameters)
        source = "fakesource"
        res = Resource.new(:type => type, :title => title, :source => source, :scope => "fakescope")

        parameters.each do |param, value|
            res.set(param, value, source)
        end
        res
    end

    def ref(type, title)
        Resource::Reference.new(:type => type, :title => title)
    end

    newtest :storage, "Stored resources" do
        hostname = "host%s" % @count
        @facts["hostname"] = hostname
        args = {:facts => @facts, :name => hostname}

        # Make all of the resources we want.  Make them at least slightly complex,
        # so we model real resources as close as we can.
        resources = []
        args[:resources] = resources
        @resources.times do |resnum|
            exec = mkresource("exec", "exec%s" % resnum, :command => "/bin/echo do something %s" % resnum)
            exec.tags = %w{exec one} << "exec%s" % resnum
            user = mkresource("user", "user%s" % resnum, :uid => resnum.to_s, :require => ref("exec", "exec%s" % resnum))
            user.tags = %w{user one two} << "user%s" % resnum
            file = mkresource("file", "/tmp/file%s" % resnum, :owner => resnum.to_s, :require => [ref("exec", "exec%s" % resnum), ref("user", "user%s" % resnum)])
            file.tags = %w{file one three} << "file%s" % resnum
            file.exported = true
            resources << exec << user << file
        end
        Puppet::Rails::Host.store(args)
    end
end

Suite.new :report, "Reports interactions" do
    def prepare
        Puppet::Transaction::Report.terminus_class = :rest
    end

    newtest :empty, "send empty report" do
        report = Puppet::Transaction::Report.new
        report.time = Time.now
        report.save
    end

    newtest :fake, "send fake report" do
        report = Puppet::Transaction::Report.new

        resourcemetrics = {
            :total => 12,
            :out_of_sync => 20,
            :applied => 45,
            :skipped => 1,
            :restarted => 23,
            :failed_restarts => 1,
            :scheduled => 10
        }
        report.newmetric(:resources, resourcemetrics)

        timemetrics = {
            :resource1 => 10,
            :resource2 => 50,
            :resource3 => 40,
            :resource4 => 20,
        }
        report.newmetric(:times, timemetrics)

        report.newmetric(:changes,
            :total => 20
        )

        report.time = Time.now
        report.save
    end
end


$cmdargs = [
    [ "--compile",  "-c", GetoptLong::NO_ARGUMENT       ],
    [ "--describe",       GetoptLong::REQUIRED_ARGUMENT ],
    [ "--retrieve", "-R", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--fork",           GetoptLong::REQUIRED_ARGUMENT ],
    [ "--fqdn",     "-F", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--suite",    "-s", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--test",     "-t", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--pause",    "-p", GetoptLong::NO_ARGUMENT       ],
    [ "--repeat",   "-r", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--node",     "-n", GetoptLong::REQUIRED_ARGUMENT ],
    [ "--debug",    "-d", GetoptLong::NO_ARGUMENT       ],
    [ "--help",     "-h", GetoptLong::NO_ARGUMENT       ],
    [ "--list",     "-l", GetoptLong::NO_ARGUMENT       ],
    [ "--verbose",  "-v", GetoptLong::NO_ARGUMENT       ],
    [ "--version",  "-V", GetoptLong::NO_ARGUMENT       ],
]

# Add all of the config parameters as valid $options.
Puppet.settings.addargs($cmdargs)
Puppet::Util::Log.newdestination(:console)

Puppet::Node.terminus_class = :plain
Puppet::Node.cache_class = :yaml
Puppet::Node::Facts.terminus_class = :facter
Puppet::Node::Facts.cache_class = :yaml

result = GetoptLong.new(*$cmdargs)

$args = {}

$options = {:repeat => 1, :fork => 0, :pause => false, :nodes => []}

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[:suite] = :configuration
                $options[:test] = :compile
            when "--retrieve"
                $options[:suite] = :file
                $options[:test] = :retrieve
                $options[:file] = arg
            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[:suite] = :file
                $options[:test] = :describe
                $options[:file] = arg
            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)
            when "--suite"
                $options[:suite] = arg.intern
            when "--test"
                $options[:test] = arg.intern
            when "--file"
                $options[:file] = arg
            when "--pause"
                $options[:pause] = true
            when "--node"
                $options[:nodes] << arg
            when "--list"
                Suite.suites.sort { |a,b| a.to_s <=> b.to_s }.each do |suite_name|
                    suite = Suite[suite_name]
                    tests = suite.tests.sort { |a,b| a.to_s <=> b.to_s }.join(", ")
                    puts "%20s: %s" % [suite_name, tests]
                end
                exit(0)
            else
                Puppet.settings.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

$options[:nodes] << Puppet.settings[:certname] if $options[:nodes].empty?

$args[:Server] = Puppet[:server]

unless $options[:test]
    $options[:suite] = :remote_catalog
    $options[:test] = :getconfig
end

unless $options[:test]
    raise "A suite was specified without a test"
end

$pids = []

Suite.run($options[:test])

if $options[:fork] > 0
    Process.waitall
end