diff --git a/Android.bp b/Android.bp
index dd9b02e..628d55b 100644
--- a/Android.bp
+++ b/Android.bp
@@ -1288,6 +1288,47 @@
     ],
 }
 
+// update_engine_unittests (type: executable)
+// ========================================================
+// Main unittest file.
+cc_test {
+    name: "update_engine_recovery_unittests",
+    defaults: [
+        "ue_defaults",
+        "update_metadata-protos_exports",
+    ],
+    cflags: [
+        "-D__ANDROID_RECOVERY__",
+    ],
+
+    static_libs: [
+        "libbrillo-test-helpers",
+        "libbase",
+        "liblog",
+        "libgmock",
+        "libchrome_test_helpers",
+        "libupdate_engine_android",
+        "update_metadata-protos",
+        "libsnapshot_cow",
+        "libdm",
+    ],
+
+    header_libs: [
+        "libstorage_literals_headers",
+    ],
+    test_suites: ["device-tests"],
+    srcs: [
+        "payload_consumer/postinstall_runner_action_recovery_unittest.cc",
+        "payload_consumer/postinstall_runner_action.cc",
+        "common/subprocess.cc",
+        "common/test_utils.cc",
+        "common/utils.cc",
+        "common/dynamic_partition_control_stub.cc",
+        "common/action_processor.cc",
+        "common/error_code_utils.cc",
+    ],
+}
+
 // Brillo update payload generation script
 // ========================================================
 sh_binary {
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index 943c3e4..b46f584 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -1363,6 +1363,7 @@
   // previous ApplyPayload() call may have requested powerwash, these
   // settings would be saved in `this->install_plan_`. Inherit that setting.
   install_plan_.powerwash_required = this->install_plan_.powerwash_required;
+  install_plan_.switch_slot_on_reboot = true;
 
   CHECK_NE(install_plan_.source_slot, UINT32_MAX);
   CHECK_NE(install_plan_.target_slot, UINT32_MAX);
diff --git a/common/fake_boot_control.h b/common/fake_boot_control.h
index 8a68501..82d4827 100644
--- a/common/fake_boot_control.h
+++ b/common/fake_boot_control.h
@@ -40,6 +40,11 @@
     dynamic_partition_control_.reset(new DynamicPartitionControlStub());
   }
 
+  void SetDynamicPartitionControl(
+      std::unique_ptr<DynamicPartitionControlInterface> dynamic_control) {
+    dynamic_partition_control_ = std::move(dynamic_control);
+  }
+
   // BootControlInterface overrides.
   unsigned int GetNumSlots() const override { return num_slots_; }
   BootControlInterface::Slot GetCurrentSlot() const override {
diff --git a/payload_consumer/postinstall_runner_action.cc b/payload_consumer/postinstall_runner_action.cc
index da9075a..3b09916 100644
--- a/payload_consumer/postinstall_runner_action.cc
+++ b/payload_consumer/postinstall_runner_action.cc
@@ -35,6 +35,7 @@
 #include "update_engine/common/action_processor.h"
 #include "update_engine/common/boot_control_interface.h"
 #include "update_engine/common/error_code_utils.h"
+#include "update_engine/common/platform_constants.h"
 #include "update_engine/common/subprocess.h"
 #include "update_engine/common/utils.h"
 
@@ -108,14 +109,17 @@
   auto dynamic_control = boot_control_->GetDynamicPartitionControl();
   CHECK(dynamic_control);
 
-  // Mount snapshot partitions for Virtual AB Compression Compression.
-  if (dynamic_control->UpdateUsesSnapshotCompression()) {
-    // If we are switching slots, then we are required to MapAllPartitions,
-    // as FinishUpdate() requires all partitions to be mapped.
-    // And switching slots requires FinishUpdate() to be called first
+  // Mount snapshot partitions for Virtual AB updates.
+  // If we are switching slots, then we are required to MapAllPartitions,
+  // as FinishUpdate() requires all partitions to be mapped.
+  // And switching slots requires FinishUpdate() to be called first
+  if (dynamic_control->GetVirtualAbFeatureFlag().IsEnabled() &&
+      !constants::kIsRecovery) {
     if (!install_plan_.partitions.empty() ||
         install_plan_.switch_slot_on_reboot) {
       if (!dynamic_control->MapAllPartitions()) {
+        LOG(ERROR) << "Failed to map all partitions, this would cause "
+                      "FinishUpdate to fail. Abort early.";
         return CompletePostinstall(ErrorCode::kPostInstallMountError);
       }
     }
diff --git a/payload_consumer/postinstall_runner_action_recovery_unittest.cc b/payload_consumer/postinstall_runner_action_recovery_unittest.cc
new file mode 100644
index 0000000..3a869c9
--- /dev/null
+++ b/payload_consumer/postinstall_runner_action_recovery_unittest.cc
@@ -0,0 +1,202 @@
+//
+// Copyright (C) 2012 The Android Open Source Project
+//
+// 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.
+//
+
+#include "gmock/gmock.h"
+#include "update_engine/payload_consumer/postinstall_runner_action.h"
+
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <memory>
+#include <string>
+#include <utility>
+#include "common/dynamic_partition_control_interface.h"
+
+#include <base/bind.h>
+#include <base/files/file_util.h>
+#include <base/message_loop/message_loop.h>
+#include <android-base/stringprintf.h>
+#include <brillo/message_loops/base_message_loop.h>
+#include <brillo/message_loops/message_loop_utils.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "update_engine/common/fake_boot_control.h"
+#include "update_engine/common/fake_hardware.h"
+#include "update_engine/common/test_utils.h"
+#include "update_engine/common/utils.h"
+#include "update_engine/common/mock_dynamic_partition_control.h"
+
+using brillo::MessageLoop;
+using chromeos_update_engine::test_utils::ScopedLoopbackDeviceBinder;
+using std::string;
+using testing::_;
+using testing::AtLeast;
+using testing::Return;
+
+namespace chromeos_update_engine {
+
+class PostinstActionProcessorDelegate : public ActionProcessorDelegate {
+ public:
+  PostinstActionProcessorDelegate() = default;
+  void ProcessingDone(const ActionProcessor* processor,
+                      ErrorCode code) override {
+    MessageLoop::current()->BreakLoop();
+    processing_done_called_ = true;
+  }
+  void ProcessingStopped(const ActionProcessor* processor) override {
+    MessageLoop::current()->BreakLoop();
+    processing_stopped_called_ = true;
+  }
+
+  void ActionCompleted(ActionProcessor* processor,
+                       AbstractAction* action,
+                       ErrorCode code) override {
+    if (action->Type() == PostinstallRunnerAction::StaticType()) {
+      code_ = code;
+      code_set_ = true;
+    }
+  }
+
+  ErrorCode code_{ErrorCode::kError};
+  bool code_set_{false};
+  bool processing_done_called_{false};
+  bool processing_stopped_called_{false};
+};
+
+class MockPostinstallRunnerActionDelegate
+    : public PostinstallRunnerAction::DelegateInterface {
+ public:
+  MOCK_METHOD1(ProgressUpdate, void(double progress));
+};
+
+class PostinstallRunnerActionTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    loop_.SetAsCurrent();
+    {
+      auto mock_dynamic_control =
+          std::make_unique<MockDynamicPartitionControl>();
+      mock_dynamic_control_ = mock_dynamic_control.get();
+      fake_boot_control_.SetDynamicPartitionControl(
+          std::move(mock_dynamic_control));
+    }
+    ON_CALL(*mock_dynamic_control_, FinishUpdate(_))
+        .WillByDefault(Return(true));
+    ON_CALL(*mock_dynamic_control_, GetVirtualAbFeatureFlag())
+        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
+  }
+
+  // Setup an action processor and run the PostinstallRunnerAction with a single
+  // partition |device_path|, running the |postinstall_program| command from
+  // there.
+  void RunPostinstallAction(bool powerwash_required, bool save_rollback_data);
+
+  void RunPostinstallActionWithInstallPlan(const InstallPlan& install_plan);
+
+ public:
+  void ResumeRunningAction() {
+    ASSERT_NE(nullptr, postinstall_action_);
+    postinstall_action_->ResumeAction();
+  }
+
+ protected:
+  base::MessageLoopForIO base_loop_;
+  brillo::BaseMessageLoop loop_{&base_loop_};
+
+  FakeBootControl fake_boot_control_;
+  FakeHardware fake_hardware_;
+  MockDynamicPartitionControl* mock_dynamic_control_;
+  PostinstActionProcessorDelegate processor_delegate_;
+
+  // The PostinstallRunnerAction delegate receiving the progress updates.
+  PostinstallRunnerAction::DelegateInterface* setup_action_delegate_{nullptr};
+
+  // A pointer to the posinstall_runner action and the processor.
+  PostinstallRunnerAction* postinstall_action_{nullptr};
+  ActionProcessor* processor_{nullptr};
+};
+
+void PostinstallRunnerActionTest::RunPostinstallAction(
+    bool powerwash_required, bool save_rollback_data) {
+  InstallPlan::Partition part;
+  part.name = "part";
+  part.target_path = "/dev/invalid";
+  part.readonly_target_path = "/dev/invalid";
+  part.run_postinstall = false;
+  part.postinstall_path.clear();
+  InstallPlan install_plan;
+  install_plan.partitions = {part};
+  install_plan.download_url = "http://127.0.0.1:8080/update";
+  install_plan.powerwash_required = powerwash_required;
+  RunPostinstallActionWithInstallPlan(install_plan);
+}
+
+void PostinstallRunnerActionTest::RunPostinstallActionWithInstallPlan(
+    const chromeos_update_engine::InstallPlan& install_plan) {
+  ActionProcessor processor;
+  processor_ = &processor;
+  auto feeder_action = std::make_unique<ObjectFeederAction<InstallPlan>>();
+  feeder_action->set_obj(install_plan);
+  auto runner_action = std::make_unique<PostinstallRunnerAction>(
+      &fake_boot_control_, &fake_hardware_);
+  postinstall_action_ = runner_action.get();
+  base::FilePath temp_dir;
+  TEST_AND_RETURN(base::CreateNewTempDirectory("postinstall", &temp_dir));
+  postinstall_action_->SetMountDir(temp_dir.value());
+  runner_action->set_delegate(setup_action_delegate_);
+  BondActions(feeder_action.get(), runner_action.get());
+  auto collector_action =
+      std::make_unique<ObjectCollectorAction<InstallPlan>>();
+  BondActions(runner_action.get(), collector_action.get());
+  processor.EnqueueAction(std::move(feeder_action));
+  processor.EnqueueAction(std::move(runner_action));
+  processor.EnqueueAction(std::move(collector_action));
+  processor.set_delegate(&processor_delegate_);
+
+  loop_.PostTask(
+      FROM_HERE,
+      base::Bind(
+          [](ActionProcessor* processor) { processor->StartProcessing(); },
+          base::Unretained(&processor)));
+  loop_.Run();
+  ASSERT_FALSE(processor.IsRunning());
+  postinstall_action_ = nullptr;
+  processor_ = nullptr;
+  ASSERT_TRUE(processor_delegate_.processing_stopped_called_ ||
+              processor_delegate_.processing_done_called_);
+  if (processor_delegate_.processing_done_called_) {
+    // Validation check that the code was set when the processor finishes.
+    ASSERT_TRUE(processor_delegate_.code_set_);
+  }
+}
+
+// Test that postinstall succeeds in the simple case of running the default
+// /postinst command which only exits 0.
+TEST_F(PostinstallRunnerActionTest, RunAsRootSimpleTest) {
+  EXPECT_CALL(*mock_dynamic_control_, GetVirtualAbFeatureFlag())
+      .WillOnce(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
+  RunPostinstallAction(false, false);
+  ASSERT_EQ(ErrorCode::kSuccess, processor_delegate_.code_);
+  ASSERT_TRUE(processor_delegate_.processing_done_called_);
+
+  // Since powerwash_required was false, this should not trigger a powerwash.
+  ASSERT_FALSE(fake_hardware_.IsPowerwashScheduled());
+  ASSERT_FALSE(fake_hardware_.GetIsRollbackPowerwashScheduled());
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/postinstall_runner_action_unittest.cc b/payload_consumer/postinstall_runner_action_unittest.cc
index 028402a..8579144 100644
--- a/payload_consumer/postinstall_runner_action_unittest.cc
+++ b/payload_consumer/postinstall_runner_action_unittest.cc
@@ -23,6 +23,7 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include "common/dynamic_partition_control_interface.h"
 
 #include <base/bind.h>
 #include <base/files/file_util.h>
@@ -44,10 +45,14 @@
 #include "update_engine/common/subprocess.h"
 #include "update_engine/common/test_utils.h"
 #include "update_engine/common/utils.h"
+#include "update_engine/common/mock_dynamic_partition_control.h"
 
 using brillo::MessageLoop;
 using chromeos_update_engine::test_utils::ScopedLoopbackDeviceBinder;
 using std::string;
+using testing::_;
+using testing::AtLeast;
+using testing::Return;
 
 namespace chromeos_update_engine {
 
@@ -95,6 +100,21 @@
     // stored in the "disk_ext2_unittest.img" image.
     postinstall_image_ =
         test_utils::GetBuildArtifactsPath("gen/disk_ext2_unittest.img");
+    {
+      auto mock_dynamic_control =
+          std::make_unique<MockDynamicPartitionControl>();
+      mock_dynamic_control_ = mock_dynamic_control.get();
+      fake_boot_control_.SetDynamicPartitionControl(
+          std::move(mock_dynamic_control));
+    }
+    ON_CALL(*mock_dynamic_control_, FinishUpdate(_))
+        .WillByDefault(Return(true));
+    ON_CALL(*mock_dynamic_control_, MapAllPartitions())
+        .WillByDefault(Return(true));
+    ON_CALL(*mock_dynamic_control_, UnmapAllPartitions())
+        .WillByDefault(Return(true));
+    ON_CALL(*mock_dynamic_control_, GetVirtualAbFeatureFlag())
+        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
   }
 
   // Setup an action processor and run the PostinstallRunnerAction with a single
@@ -173,6 +193,7 @@
 
   FakeBootControl fake_boot_control_;
   FakeHardware fake_hardware_;
+  MockDynamicPartitionControl* mock_dynamic_control_;
   PostinstActionProcessorDelegate processor_delegate_;
 
   // The PostinstallRunnerAction delegate receiving the progress updates.
@@ -398,6 +419,7 @@
 // Test that we can cancel a postinstall action while it is running.
 TEST_F(PostinstallRunnerActionTest, RunAsRootCancelPostinstallActionTest) {
   ScopedLoopbackDeviceBinder loop(postinstall_image_, false, nullptr);
+  EXPECT_CALL(*mock_dynamic_control_, MapAllPartitions()).Times(AtLeast(1));
 
   // Wait for the action to start and then cancel it.
   CancelWhenStarted();
@@ -411,6 +433,10 @@
 // Test that we parse and process the progress reports from the progress
 // file descriptor.
 TEST_F(PostinstallRunnerActionTest, RunAsRootProgressUpdatesTest) {
+  EXPECT_CALL(*mock_dynamic_control_, MapAllPartitions())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(true));
+  EXPECT_CALL(*mock_dynamic_control_, FinishUpdate(_)).Times(AtLeast(1));
   testing::StrictMock<MockPostinstallRunnerActionDelegate> mock_delegate_;
   testing::InSequence s;
   EXPECT_CALL(mock_delegate_, ProgressUpdate(0));