From 0d918c46d753d8008b98bec4682f833abb6d7029 Mon Sep 17 00:00:00 2001
From: Tony Vu <tony.vu@puppetlabs.com>
Date: Thu, 12 May 2016 18:20:51 -0700
Subject: [PATCH] (BKR-656) refactor pe_ver setting into independent method

Previous to this commit, transforming a host object prior to upgrading
was handled in the upgrade_pe_on method. This change removes that logic
from that method and allows for independent transformation to happen in
a new prep_host_for_upgrade method.
---
 lib/beaker-pe/install/pe_utils.rb       | 35 ++++++++++++++++-------
 spec/beaker-pe/install/pe_utils_spec.rb | 38 +++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/lib/beaker-pe/install/pe_utils.rb b/lib/beaker-pe/install/pe_utils.rb
index ea197b00..e368fd3e 100644
--- a/lib/beaker-pe/install/pe_utils.rb
+++ b/lib/beaker-pe/install/pe_utils.rb
@@ -634,23 +634,36 @@ def upgrade_pe_on upgrade_hosts, opts, path=nil
             end
             # get new version information
             hosts.each do |host|
-              host['pe_dir'] = host['pe_upgrade_dir'] || path
-              if host['platform'] =~ /windows/
-                host['pe_ver'] = host['pe_upgrade_ver'] || opts['pe_upgrade_ver'] ||
-                  Options::PEVersionScraper.load_pe_version(host['pe_dir'], opts[:pe_version_file_win])
-              else
-                host['pe_ver'] = host['pe_upgrade_ver'] || opts['pe_upgrade_ver'] ||
-                  Options::PEVersionScraper.load_pe_version(host['pe_dir'], opts[:pe_version_file])
-              end
-              if version_is_less(host['pe_ver'], '3.0')
-                host['pe_installer'] ||= 'puppet-enterprise-upgrader'
-              end
+              prep_host_for_upgrade(host, opts, path)
             end
             do_install(sorted_hosts, opts.merge({:type => :upgrade, :set_console_password => set_console_password}))
             opts['upgrade'] = true
           end
         end
 
+        #Prep a host object for upgrade; used inside upgrade_pe_on
+        # @param [Host] host A single host object to prepare for upgrade
+        # !macro common_opts
+        # @param [String] path A path (either local directory or a URL to a listing of PE builds).
+        #                      Will contain a LATEST file indicating the latest build to install.
+        #                      This is ignored if a pe_upgrade_ver and pe_upgrade_dir are specified
+        #                      in the host configuration file.
+        # @example
+        #  pre_host_for_upgrade(master, {}, "http://neptune.puppetlabs.lan/3.0/ci-ready/")
+        def prep_host_for_upgrade(host, opts={}, path='')
+          host['pe_dir'] = host['pe_upgrade_dir'] || path
+          if host['platform'] =~ /windows/
+            host['pe_ver'] = host['pe_upgrade_ver'] || opts['pe_upgrade_ver'] ||
+              Options::PEVersionScraper.load_pe_version(host['pe_dir'], opts[:pe_version_file_win])
+          else
+            host['pe_ver'] = host['pe_upgrade_ver'] || opts['pe_upgrade_ver'] ||
+              Options::PEVersionScraper.load_pe_version(host['pe_dir'], opts[:pe_version_file])
+          end
+          if version_is_less(host['pe_ver'], '3.0')
+            host['pe_installer'] ||= 'puppet-enterprise-upgrader'
+          end
+        end
+
         #Create the Higgs install command string based upon the host and options settings.  Installation command will be run as a
         #background process.  The output of the command will be stored in the provided host['higgs_file'].
         # @param [Host] host The host that Higgs is to be installed on
diff --git a/spec/beaker-pe/install/pe_utils_spec.rb b/spec/beaker-pe/install/pe_utils_spec.rb
index 5da5ee46..f2e7be6b 100644
--- a/spec/beaker-pe/install/pe_utils_spec.rb
+++ b/spec/beaker-pe/install/pe_utils_spec.rb
@@ -52,6 +52,44 @@ def logger
                                                 :type => 'pe',
                                                 :working_dir => '/tmp',
                                                 :dist => 'puppet-enterprise-3.7.1-rc0-78-gffc958f-eos-4-i386' } ) }
+
+  context '#prep_host_for_upgrade' do
+
+    it 'sets per host options before global options' do
+      opts['pe_upgrade_ver'] = 'options-specific-var'
+      hosts.each do |host|
+        host['pe_upgrade_dir'] = 'host-specific-pe-dir'
+        host['pe_upgrade_ver'] = 'host-specific-pe-ver'
+        subject.prep_host_for_upgrade(host, opts, 'argument-specific-pe-dir')
+        expect(host['pe_dir']).to eq('host-specific-pe-dir')
+        expect(host['pe_ver']).to eq('host-specific-pe-ver')
+      end
+    end
+
+    it 'sets global options when no host options are available' do
+      opts['pe_upgrade_ver'] = 'options-specific-var'
+      hosts.each do |host|
+        host['pe_upgrade_dir'] = nil
+        host['pe_upgrade_ver'] = nil
+        subject.prep_host_for_upgrade(host, opts, 'argument-specific-pe-dir')
+        expect(host['pe_dir']).to eq('argument-specific-pe-dir')
+        expect(host['pe_ver']).to eq('options-specific-var')
+      end
+    end
+
+    it 'calls #load_pe_version when neither global or host options are present' do
+      opts['pe_upgrade_ver'] = nil
+      hosts.each do |host|
+        host['pe_upgrade_dir'] = nil
+        host['pe_upgrade_ver'] = nil
+        expect( Beaker::Options::PEVersionScraper ).to receive(:load_pe_version).and_return('file_version')
+        subject.prep_host_for_upgrade(host, opts, 'argument-specific-pe-dir')
+        expect(host['pe_ver']).to eq('file_version')
+        expect(host['pe_dir']).to eq('argument-specific-pe-dir')
+      end
+    end
+  end
+
   context '#configure_pe_defaults_on' do
     it 'uses aio paths for hosts of role aio' do
       hosts.each do |host|