Map all partitions before calling FinishUpdate for VAB as well

In virtual AB updates, all partitions must be mapped before calling
FinishUpdate. (same requirement as VABC).

Test: install VAB update w/o slot switch, then perform slot switch
Bug: 393082101
Change-Id: Iec7965b2d1a75b8a61507d75c83c733bfac38ce4
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..1717dd7 100644
--- a/payload_consumer/postinstall_runner_action.cc
+++ b/payload_consumer/postinstall_runner_action.cc
@@ -108,16 +108,15 @@
   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
-    if (!install_plan_.partitions.empty() ||
-        install_plan_.switch_slot_on_reboot) {
-      if (!dynamic_control->MapAllPartitions()) {
-        return CompletePostinstall(ErrorCode::kPostInstallMountError);
-      }
+  // 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 (!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);
     }
   }
 
@@ -324,10 +323,10 @@
 
 void PostinstallRunnerAction::OnProgressFdReady() {
   char buf[1024];
-  size_t bytes_read;
+  size_t bytes_read{};
   do {
     bytes_read = 0;
-    bool eof;
+    bool eof = false;
     bool ok =
         utils::ReadAll(progress_fd_, buf, std::size(buf), &bytes_read, &eof);
     progress_buffer_.append(buf, bytes_read);
diff --git a/payload_consumer/postinstall_runner_action_unittest.cc b/payload_consumer/postinstall_runner_action_unittest.cc
index 028402a..7416e66 100644
--- a/payload_consumer/postinstall_runner_action_unittest.cc
+++ b/payload_consumer/postinstall_runner_action_unittest.cc
@@ -44,10 +44,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 +99,19 @@
     // 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));
   }
 
   // Setup an action processor and run the PostinstallRunnerAction with a single
@@ -173,6 +190,7 @@
 
   FakeBootControl fake_boot_control_;
   FakeHardware fake_hardware_;
+  MockDynamicPartitionControl* mock_dynamic_control_;
   PostinstActionProcessorDelegate processor_delegate_;
 
   // The PostinstallRunnerAction delegate receiving the progress updates.
@@ -398,6 +416,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 +430,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));