diff --git a/.sync.yml b/.sync.yml
index b2375f02..e97335cf 100644
--- a/.sync.yml
+++ b/.sync.yml
@@ -14,6 +14,8 @@ Rakefile:
   changelog_since_tag: '2.1.0'
   extras:
     - 'PuppetSyntax.exclude_paths = ["plans/**/*.pp", "vendor/**/*"]'
+spec/spec_helper.rb:
+    mock_with: ':rspec'
 .gitignore:
   paths:
     - '.rerun.json'
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6a65d52b..c5966f3b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,7 +5,7 @@
 
 ### Improvements
 
-- Handle exit code 11 from replica upgrade (still waiting for PuppetDB sync to complete) gracefully.
+- Handle exit code 11 from replica upgrade task gracefully. Code 11 means "PuppetDB sync in progress but not yet complete".
 
 ## 2.4.5
 ### Summary
diff --git a/metadata.json b/metadata.json
index 29b82361..14aecacf 100644
--- a/metadata.json
+++ b/metadata.json
@@ -62,7 +62,7 @@
       "version_requirement": ">= 6.0.2 < 7.0.0"
     }
   ],
-  "pdk-version": "1.18.0",
+  "pdk-version": "1.18.1",
   "template-url": "https://github.com/puppetlabs/pdk-templates.git#1.18.0",
   "template-ref": "tags/1.18.0-0-g095317c"
 }
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index d3778cac..b367fded 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -1,5 +1,9 @@
 # frozen_string_literal: true
 
+RSpec.configure do |c|
+  c.mock_with :rspec
+end
+
 require 'puppetlabs_spec_helper/module_spec_helper'
 require 'rspec-puppet-facts'
 
diff --git a/spec/spec_helper_local.rb b/spec/spec_helper_local.rb
index 3622c6b4..5f9bc2da 100644
--- a/spec/spec_helper_local.rb
+++ b/spec/spec_helper_local.rb
@@ -2,6 +2,10 @@
 
 require 'puppet'
 
+# This environment variable can be read by Ruby Bolt tasks to prevent unwanted
+# auto-execution, enabling easy unit testing.
+ENV["RSPEC_UNIT_TEST_MODE"] ||= "TRUE"
+
 if Gem::Version.new(Puppet.version) >= Gem::Version.new('6.0.0')
   begin
     require 'bolt_spec/plans'
diff --git a/spec/unit/task/puppet_infra_upgrade_spec.rb b/spec/unit/task/puppet_infra_upgrade_spec.rb
new file mode 100644
index 00000000..6a8e924f
--- /dev/null
+++ b/spec/unit/task/puppet_infra_upgrade_spec.rb
@@ -0,0 +1,43 @@
+require 'spec_helper'
+require_relative '../../../tasks/puppet_infra_upgrade'
+
+describe PEAdm::Task::PuppetInfraUpgrade do
+  context 'replica' do
+    let(:upgrade) do
+      described_class.new('type' => 'replica',
+                          'targets' => ['replica.example.com'],
+                          'wait_until_connected_timeout' => 120)
+    end
+
+    it 'Returns when the command exits with an expected code' do
+      status_dbl_0 = double('Status', :exitstatus => 0)
+      status_dbl_11 = double('Status', :exitstatus => 11)
+      allow(STDOUT).to receive(:puts)
+      allow(upgrade).to receive(:wait_until_connected)
+
+      allow(Open3).to receive(:capture2e).and_return(['hello world', status_dbl_0])
+      expect(upgrade.execute!).to eq(nil)
+
+      allow(Open3).to receive(:capture2e).and_return(['hello world', status_dbl_11])
+      expect(upgrade.execute!).to eq(nil)
+    end
+
+    it 'Exits non-zero when the command exits with an unexpected code' do
+      status_dbl_1 = double('Status', :exitstatus => 1)
+      allow(STDOUT).to receive(:puts)
+      allow(upgrade).to receive(:wait_until_connected)
+      allow(Open3).to receive(:capture2e).and_return(['hello world', status_dbl_1])
+      expect { upgrade.execute! }.to raise_error(SystemExit) do |error|
+        expect(error.status).to eq(1)
+      end
+    end
+  end
+
+  context 'compiler' do
+    let(:upgrade) do
+      described_class.new('type' => 'compiler',
+                          'targets' => ['replica.example.com'],
+                          'wait_until_connected_timeout' => 120)
+    end
+  end
+end
diff --git a/tasks/puppet_infra_upgrade.rb b/tasks/puppet_infra_upgrade.rb
index 0cde0884..3da754d8 100755
--- a/tasks/puppet_infra_upgrade.rb
+++ b/tasks/puppet_infra_upgrade.rb
@@ -8,76 +8,93 @@
 require 'timeout'
 require 'etc'
 
-def main
-  params     = JSON.parse(STDIN.read)
-  type       = params['type']
-  targets    = params['targets']
-  timeout    = params['wait_until_connected_timeout']
-  token_file = params['token_file'] || File.join(Etc.getpwuid.dir, '.puppetlabs', 'token')
+class PEAdm
+  class Task
+    class PuppetInfraUpgrade
+      def initialize(params)
+        @type       = params['type']
+        @targets    = params['targets']
+        @timeout    = params['wait_until_connected_timeout']
+        @token_file = params['token_file']
+      end
 
-  exit 0 if targets.empty?
+      def execute!
+        exit 0 if @targets.empty?
+        token_file = @token_file || File.join(Etc.getpwuid.dir, '.puppetlabs', 'token')
 
-  cmd = ['/opt/puppetlabs/bin/puppet-infrastructure', '--render-as', 'json', 'upgrade']
-  cmd << '--token-file' << token_file unless params['token_file'].nil?
-  cmd << type << targets.join(',')
+        cmd = ['/opt/puppetlabs/bin/puppet-infrastructure', '--render-as', 'json', 'upgrade']
+        cmd << '--token-file' << token_file unless @token_file.nil?
+        cmd << @type << @targets.join(',')
 
-  wait_until_connected(nodes: targets, token_file: token_file, timeout: timeout)
+        wait_until_connected(nodes: @targets, token_file: token_file, timeout: @timeout)
 
-  stdouterr, status = Open3.capture2e(*cmd)
-  puts stdouterr
-  if status.success?
-    exit 0
-  elsif status.exitstatus == 11 # Waiting for PuppetDB sync to complete, but otherwise successful
-    exit 0
-  else
-    exit status.exitstatus
-  end
-end
+        stdouterr, status = Open3.capture2e(*cmd)
+        STDOUT.puts stdouterr
 
-def inventory_uri
-  @inventory_uri ||= URI.parse('https://localhost:8143/orchestrator/v1/inventory')
-end
+        # Exit code 11 indicates PuppetDB sync in progress, just not yet
+        # finished. We consider that success.
+        if [0, 11].include?(status.exitstatus)
+          return
+        else
+          exit status.exitstatus
+        end
+      end
 
-def request_object(nodes:, token_file:)
-  token = File.read(token_file)
-  body = {
-    'nodes' => nodes,
-  }.to_json
+      def inventory_uri
+        @inventory_uri ||= URI.parse('https://localhost:8143/orchestrator/v1/inventory')
+      end
 
-  request = Net::HTTP::Post.new(inventory_uri.request_uri)
-  request['Content-Type'] = 'application/json'
-  request['X-Authentication'] = token
-  request.body = body
+      def request_object(nodes:, token_file:)
+        token = File.read(token_file)
+        body = {
+          'nodes' => nodes,
+        }.to_json
 
-  request
-end
+        request = Net::HTTP::Post.new(inventory_uri.request_uri)
+        request['Content-Type'] = 'application/json'
+        request['X-Authentication'] = token
+        request.body = body
 
-def http_object
-  http = Net::HTTP.new(inventory_uri.host, inventory_uri.port)
-  http.use_ssl = true
-  http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+        request
+      end
 
-  http
-end
+      def http_object
+        http = Net::HTTP.new(inventory_uri.host, inventory_uri.port)
+        http.use_ssl = true
+        http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+
+        http
+      end
 
-def wait_until_connected(nodes:, token_file:, timeout: 120)
-  http = http_object
-  request = request_object(nodes: nodes, token_file: token_file)
-  inventory = {}
-  Timeout.timeout(timeout) do
-    loop do
-      response = http.request(request)
-      raise unless response.is_a? Net::HTTPSuccess
-      inventory = JSON.parse(response.body)
-      break if inventory['items'].all? { |item| item['connected'] }
-      sleep(1)
+      def wait_until_connected(nodes:, token_file:, timeout: 120)
+        http = http_object
+        request = request_object(nodes: nodes, token_file: token_file)
+        inventory = {}
+        Timeout.timeout(timeout) do
+          loop do
+            response = http.request(request)
+            unless response.is_a? Net::HTTPSuccess
+              raise "Unexpected result from orchestrator: #{response.class}\n#{response}"
+            end
+            inventory = JSON.parse(response.body)
+            break if inventory['items'].all? { |item| item['connected'] }
+            sleep(1)
+          end
+        end
+      rescue Timeout::Error
+        raise 'Timed out waiting for nodes to be connected to orchestrator: ' +
+              inventory['items'].reject { |item| item['connected'] }
+                                .map { |item| item['name'] }
+                                .to_s
+      end
     end
   end
-rescue Timeout::Error
-  raise 'Timed out waiting for nodes to be connected to orchestrator: ' +
-        inventory['items'].reject { |item| item['connected'] }
-                          .map { |item| item['name'] }
-                          .to_s
 end
 
-main
+# Run the task unless an environment flag has been set, signaling not to. The
+# environment flag is used to disable auto-execution and enable Ruby unit
+# testing of this task.
+unless ENV['RSPEC_UNIT_TEST_MODE']
+  upgrade = PEAdm::Task::PuppetInfraUpgrade.new(JSON.parse(STDIN.read))
+  upgrade.execute!
+end