Report retrofit and COW image size am: 4ed0512acc

Change-Id: If15f3f192bede7d5aac20987d8c3d1491452ab43
diff --git a/METADATA b/METADATA
new file mode 100644
index 0000000..d97975c
--- /dev/null
+++ b/METADATA
@@ -0,0 +1,3 @@
+third_party {
+  license_type: NOTICE
+}
diff --git a/cleanup_previous_update_action.cc b/cleanup_previous_update_action.cc
index 88dbc57..1a2476f 100644
--- a/cleanup_previous_update_action.cc
+++ b/cleanup_previous_update_action.cc
@@ -31,7 +31,7 @@
 #include "update_engine/payload_consumer/delta_performer.h"
 
 using android::base::GetBoolProperty;
-using android::snapshot::SnapshotManager;
+using android::snapshot::ISnapshotManager;
 using android::snapshot::SnapshotMergeStats;
 using android::snapshot::UpdateState;
 using brillo::MessageLoop;
@@ -56,7 +56,7 @@
 CleanupPreviousUpdateAction::CleanupPreviousUpdateAction(
     PrefsInterface* prefs,
     BootControlInterface* boot_control,
-    android::snapshot::SnapshotManager* snapshot,
+    android::snapshot::ISnapshotManager* snapshot,
     CleanupPreviousUpdateActionDelegateInterface* delegate)
     : prefs_(prefs),
       boot_control_(boot_control),
@@ -65,7 +65,7 @@
       running_(false),
       cancel_failed_(false),
       last_percentage_(0),
-      merge_stats_(SnapshotMergeStats::GetInstance(*snapshot)) {}
+      merge_stats_(nullptr) {}
 
 void CleanupPreviousUpdateAction::PerformAction() {
   ResumeAction();
@@ -111,8 +111,10 @@
     processor_->ActionComplete(this, ErrorCode::kSuccess);
     return;
   }
-  // SnapshotManager is only available on VAB devices.
-  CHECK(snapshot_);
+  // SnapshotManager must be available on VAB devices.
+  CHECK(snapshot_ != nullptr);
+  merge_stats_ = snapshot_->GetSnapshotMergeStatsInstance();
+  CHECK(merge_stats_ != nullptr);
   WaitBootCompletedOrSchedule();
 }
 
diff --git a/cleanup_previous_update_action.h b/cleanup_previous_update_action.h
index 91e08b0..6f6ce07 100644
--- a/cleanup_previous_update_action.h
+++ b/cleanup_previous_update_action.h
@@ -49,7 +49,7 @@
   CleanupPreviousUpdateAction(
       PrefsInterface* prefs,
       BootControlInterface* boot_control,
-      android::snapshot::SnapshotManager* snapshot,
+      android::snapshot::ISnapshotManager* snapshot,
       CleanupPreviousUpdateActionDelegateInterface* delegate);
 
   void PerformAction() override;
@@ -67,13 +67,13 @@
  private:
   PrefsInterface* prefs_;
   BootControlInterface* boot_control_;
-  android::snapshot::SnapshotManager* snapshot_;
+  android::snapshot::ISnapshotManager* snapshot_;
   CleanupPreviousUpdateActionDelegateInterface* delegate_;
   std::unique_ptr<android::snapshot::AutoDevice> metadata_device_;
   bool running_{false};
   bool cancel_failed_{false};
   unsigned int last_percentage_{0};
-  android::snapshot::SnapshotMergeStats* merge_stats_;
+  android::snapshot::ISnapshotMergeStats* merge_stats_;
 
   void StartActionInternal();
   void ScheduleWaitBootCompleted();
diff --git a/dynamic_partition_control_android.cc b/dynamic_partition_control_android.cc
index 5ed604a..a9c2bb3 100644
--- a/dynamic_partition_control_android.cc
+++ b/dynamic_partition_control_android.cc
@@ -35,6 +35,7 @@
 #include <libavb/libavb.h>
 #include <libdm/dm.h>
 #include <libsnapshot/snapshot.h>
+#include <libsnapshot/snapshot_stub.h>
 
 #include "update_engine/cleanup_previous_update_action.h"
 #include "update_engine/common/boot_control_interface.h"
@@ -58,6 +59,7 @@
 using android::snapshot::OptimizeSourceCopyOperation;
 using android::snapshot::Return;
 using android::snapshot::SnapshotManager;
+using android::snapshot::SnapshotManagerStub;
 using android::snapshot::UpdateState;
 
 namespace chromeos_update_engine {
@@ -108,8 +110,10 @@
       virtual_ab_(GetFeatureFlag(kVirtualAbEnabled, kVirtualAbRetrofit)) {
   if (GetVirtualAbFeatureFlag().IsEnabled()) {
     snapshot_ = SnapshotManager::New();
-    CHECK(snapshot_ != nullptr) << "Cannot initialize SnapshotManager.";
+  } else {
+    snapshot_ = SnapshotManagerStub::New();
   }
+  CHECK(snapshot_ != nullptr) << "Cannot initialize SnapshotManager.";
 }
 
 FeatureFlag DynamicPartitionControlAndroid::GetDynamicPartitionsFeatureFlag() {
@@ -489,8 +493,13 @@
     }
   }
 
-  return PrepareDynamicPartitionsForUpdate(
-      source_slot, target_slot, manifest, delete_source);
+  TEST_AND_RETURN_FALSE(PrepareDynamicPartitionsForUpdate(
+      source_slot, target_slot, manifest, delete_source));
+
+  if (required_size != nullptr) {
+    *required_size = 0;
+  }
+  return true;
 }
 
 namespace {
diff --git a/dynamic_partition_control_android.h b/dynamic_partition_control_android.h
index 8ad7593..18a05fb 100644
--- a/dynamic_partition_control_android.h
+++ b/dynamic_partition_control_android.h
@@ -173,8 +173,19 @@
   virtual bool EraseSystemOtherAvbFooter(uint32_t source_slot,
                                          uint32_t target_slot);
 
+  // Helper for PreparePartitionsForUpdate. Used for devices with dynamic
+  // partitions updating without snapshots.
+  // If |delete_source| is set, source partitions are deleted before resizing
+  // target partitions (using DeleteSourcePartitions).
+  virtual bool PrepareDynamicPartitionsForUpdate(
+      uint32_t source_slot,
+      uint32_t target_slot,
+      const DeltaArchiveManifest& manifest,
+      bool delete_source);
+
  private:
   friend class DynamicPartitionControlAndroidTest;
+  friend class SnapshotPartitionTestP;
 
   void UnmapAllPartitions();
   bool MapPartitionInternal(const std::string& super_device,
@@ -189,15 +200,6 @@
                                uint32_t target_slot,
                                const DeltaArchiveManifest& manifest);
 
-  // Helper for PreparePartitionsForUpdate. Used for devices with dynamic
-  // partitions updating without snapshots.
-  // If |delete_source| is set, source partitions are deleted before resizing
-  // target partitions (using DeleteSourcePartitions).
-  bool PrepareDynamicPartitionsForUpdate(uint32_t source_slot,
-                                         uint32_t target_slot,
-                                         const DeltaArchiveManifest& manifest,
-                                         bool delete_source);
-
   // Helper for PreparePartitionsForUpdate. Used for snapshotted partitions for
   // Virtual A/B update.
   bool PrepareSnapshotPartitionsForUpdate(uint32_t source_slot,
@@ -258,7 +260,7 @@
   std::set<std::string> mapped_devices_;
   const FeatureFlag dynamic_partitions_;
   const FeatureFlag virtual_ab_;
-  std::unique_ptr<android::snapshot::SnapshotManager> snapshot_;
+  std::unique_ptr<android::snapshot::ISnapshotManager> snapshot_;
   std::unique_ptr<android::snapshot::AutoDevice> metadata_device_;
   bool target_supports_snapshot_ = false;
   // Whether the target partitions should be loaded as dynamic partitions. Set
diff --git a/dynamic_partition_control_android_unittest.cc b/dynamic_partition_control_android_unittest.cc
index 2081918..3738170 100644
--- a/dynamic_partition_control_android_unittest.cc
+++ b/dynamic_partition_control_android_unittest.cc
@@ -24,6 +24,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <libavb/libavb.h>
+#include <libsnapshot/mock_snapshot.h>
 
 #include "update_engine/common/mock_prefs.h"
 #include "update_engine/common/test_utils.h"
@@ -31,6 +32,7 @@
 #include "update_engine/mock_dynamic_partition_control.h"
 
 using android::dm::DmDeviceState;
+using android::snapshot::MockSnapshotManager;
 using chromeos_update_engine::test_utils::ScopedLoopbackDeviceBinder;
 using chromeos_update_engine::test_utils::ScopedTempFile;
 using std::string;
@@ -72,6 +74,17 @@
 
     ON_CALL(dynamicControl(), EraseSystemOtherAvbFooter(_, _))
         .WillByDefault(Return(true));
+
+    ON_CALL(dynamicControl(), IsRecovery()).WillByDefault(Return(false));
+
+    ON_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
+        .WillByDefault(Invoke([&](uint32_t source_slot,
+                                  uint32_t target_slot,
+                                  const DeltaArchiveManifest& manifest,
+                                  bool delete_source) {
+          return dynamicControl().RealPrepareDynamicPartitionsForUpdate(
+              source_slot, target_slot, manifest, delete_source);
+        }));
   }
 
   // Return the mocked DynamicPartitionControlInterface.
@@ -892,4 +905,112 @@
   ASSERT_EQ(new_expected, device_content);
 }
 
+class FakeAutoDevice : public android::snapshot::AutoDevice {
+ public:
+  FakeAutoDevice() : AutoDevice("") {}
+};
+
+class SnapshotPartitionTestP : public DynamicPartitionControlAndroidTestP {
+ public:
+  void SetUp() override {
+    DynamicPartitionControlAndroidTestP::SetUp();
+    ON_CALL(dynamicControl(), GetVirtualAbFeatureFlag())
+        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
+
+    snapshot_ = new NiceMock<MockSnapshotManager>();
+    dynamicControl().snapshot_.reset(snapshot_);  // takes ownership
+    EXPECT_CALL(*snapshot_, BeginUpdate()).WillOnce(Return(true));
+    EXPECT_CALL(*snapshot_, EnsureMetadataMounted())
+        .WillRepeatedly(
+            Invoke([]() { return std::make_unique<FakeAutoDevice>(); }));
+
+    manifest_ =
+        PartitionSizesToManifest({{"system", 3_GiB}, {"vendor", 1_GiB}});
+  }
+  void ExpectCreateUpdateSnapshots(android::snapshot::Return val) {
+    manifest_.mutable_dynamic_partition_metadata()->set_snapshot_enabled(true);
+    EXPECT_CALL(*snapshot_, CreateUpdateSnapshots(_))
+        .WillRepeatedly(Invoke([&, val](const auto& manifest) {
+          // Deep comparison requires full protobuf library. Comparing the
+          // pointers are sufficient.
+          EXPECT_EQ(&manifest_, &manifest);
+          LOG(WARNING) << "CreateUpdateSnapshots returning " << val.string();
+          return val;
+        }));
+  }
+  bool PreparePartitionsForUpdate(uint64_t* required_size) {
+    return dynamicControl().PreparePartitionsForUpdate(
+        source(), target(), manifest_, true /* update */, required_size);
+  }
+  MockSnapshotManager* snapshot_ = nullptr;
+  DeltaArchiveManifest manifest_;
+};
+
+// Test happy path of PreparePartitionsForUpdate on a Virtual A/B device.
+TEST_P(SnapshotPartitionTestP, PreparePartitions) {
+  ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
+  uint64_t required_size = 0;
+  EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
+  EXPECT_EQ(0u, required_size);
+}
+
+// Test that if not enough space, required size returned by SnapshotManager is
+// passed up.
+TEST_P(SnapshotPartitionTestP, PreparePartitionsNoSpace) {
+  ExpectCreateUpdateSnapshots(android::snapshot::Return::NoSpace(1_GiB));
+  uint64_t required_size = 0;
+  EXPECT_FALSE(PreparePartitionsForUpdate(&required_size));
+  EXPECT_EQ(1_GiB, required_size);
+}
+
+// Test that in recovery, use empty space in super partition for a snapshot
+// update first.
+TEST_P(SnapshotPartitionTestP, RecoveryUseSuperEmpty) {
+  ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
+  EXPECT_CALL(dynamicControl(), IsRecovery()).WillRepeatedly(Return(true));
+  // Must not call PrepareDynamicPartitionsForUpdate if
+  // PrepareSnapshotPartitionsForUpdate succeeds.
+  EXPECT_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
+      .Times(0);
+  uint64_t required_size = 0;
+  EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
+  EXPECT_EQ(0u, required_size);
+}
+
+// Test that in recovery, if CreateUpdateSnapshots throws an error, try
+// the flashing path for full updates.
+TEST_P(SnapshotPartitionTestP, RecoveryErrorShouldDeleteSource) {
+  // Expectation on PreparePartitionsForUpdate
+  ExpectCreateUpdateSnapshots(android::snapshot::Return::NoSpace(1_GiB));
+  EXPECT_CALL(dynamicControl(), IsRecovery()).WillRepeatedly(Return(true));
+  EXPECT_CALL(*snapshot_, CancelUpdate()).WillOnce(Return(true));
+  EXPECT_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))
+      .WillRepeatedly(Invoke([&](auto source_slot,
+                                 auto target_slot,
+                                 const auto& manifest,
+                                 auto delete_source) {
+        EXPECT_EQ(source(), source_slot);
+        EXPECT_EQ(target(), target_slot);
+        // Deep comparison requires full protobuf library. Comparing the
+        // pointers are sufficient.
+        EXPECT_EQ(&manifest_, &manifest);
+        EXPECT_TRUE(delete_source);
+        return dynamicControl().RealPrepareDynamicPartitionsForUpdate(
+            source_slot, target_slot, manifest, delete_source);
+      }));
+  // Expectation on PrepareDynamicPartitionsForUpdate
+  SetMetadata(source(), {{S("system"), 2_GiB}, {S("vendor"), 1_GiB}});
+  ExpectUnmap({T("system"), T("vendor")});
+  // Expect that the source partitions aren't present in target super metadata.
+  ExpectStoreMetadata({{T("system"), 3_GiB}, {T("vendor"), 1_GiB}});
+
+  uint64_t required_size = 0;
+  EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
+  EXPECT_EQ(0u, required_size);
+}
+
+INSTANTIATE_TEST_CASE_P(DynamicPartitionControlAndroidTest,
+                        SnapshotPartitionTestP,
+                        testing::Values(TestParam{0, 1}, TestParam{1, 0}));
+
 }  // namespace chromeos_update_engine
diff --git a/mock_dynamic_partition_control.h b/mock_dynamic_partition_control.h
index 1e4e5fd..1aaebd8 100644
--- a/mock_dynamic_partition_control.h
+++ b/mock_dynamic_partition_control.h
@@ -28,60 +28,55 @@
 
 namespace chromeos_update_engine {
 
-class MockDynamicPartitionControl : public DynamicPartitionControlInterface {
- public:
-  MOCK_METHOD5(MapPartitionOnDeviceMapper,
-               bool(const std::string&,
-                    const std::string&,
-                    uint32_t,
-                    bool,
-                    std::string*));
-  MOCK_METHOD0(Cleanup, void());
-  MOCK_METHOD0(GetDynamicPartitionsFeatureFlag, FeatureFlag());
-  MOCK_METHOD5(
-      PreparePartitionsForUpdate,
-      bool(uint32_t, uint32_t, const DeltaArchiveManifest&, bool, uint64_t*));
-  MOCK_METHOD0(GetVirtualAbFeatureFlag, FeatureFlag());
-  MOCK_METHOD1(FinishUpdate, bool(bool));
-  MOCK_METHOD0(CleanupSuccessfulUpdate, ErrorCode());
-  MOCK_METHOD3(GetCleanupPreviousUpdateAction,
-               std::unique_ptr<AbstractAction>(
-                   BootControlInterface*,
-                   PrefsInterface*,
-                   CleanupPreviousUpdateActionDelegateInterface*));
-};
-
 class MockDynamicPartitionControlAndroid
     : public DynamicPartitionControlAndroid {
  public:
-  MOCK_METHOD5(MapPartitionOnDeviceMapper,
-               bool(const std::string&,
-                    const std::string&,
-                    uint32_t,
-                    bool,
-                    std::string*));
-  MOCK_METHOD1(UnmapPartitionOnDeviceMapper, bool(const std::string&));
-  MOCK_METHOD0(Cleanup, void());
-  MOCK_METHOD1(DeviceExists, bool(const std::string&));
-  MOCK_METHOD1(GetState, ::android::dm::DmDeviceState(const std::string&));
-  MOCK_METHOD2(GetDmDevicePathByName, bool(const std::string&, std::string*));
-  MOCK_METHOD3(LoadMetadataBuilder,
-               std::unique_ptr<::android::fs_mgr::MetadataBuilder>(
-                   const std::string&, uint32_t, uint32_t));
-  MOCK_METHOD3(StoreMetadata,
-               bool(const std::string&,
-                    android::fs_mgr::MetadataBuilder*,
-                    uint32_t));
-  MOCK_METHOD1(GetDeviceDir, bool(std::string*));
-  MOCK_METHOD0(GetDynamicPartitionsFeatureFlag, FeatureFlag());
-  MOCK_METHOD1(GetSuperPartitionName, std::string(uint32_t));
-  MOCK_METHOD0(GetVirtualAbFeatureFlag, FeatureFlag());
-  MOCK_METHOD1(FinishUpdate, bool(bool));
-  MOCK_METHOD5(
-      GetSystemOtherPath,
-      bool(uint32_t, uint32_t, const std::string&, std::string*, bool*));
-  MOCK_METHOD2(EraseSystemOtherAvbFooter, bool(uint32_t, uint32_t));
-  MOCK_METHOD0(IsAvbEnabledOnSystemOther, std::optional<bool>());
+  MOCK_METHOD(
+      bool,
+      MapPartitionOnDeviceMapper,
+      (const std::string&, const std::string&, uint32_t, bool, std::string*),
+      (override));
+  MOCK_METHOD(bool,
+              UnmapPartitionOnDeviceMapper,
+              (const std::string&),
+              (override));
+  MOCK_METHOD(void, Cleanup, (), (override));
+  MOCK_METHOD(bool, DeviceExists, (const std::string&), (override));
+  MOCK_METHOD(::android::dm::DmDeviceState,
+              GetState,
+              (const std::string&),
+              (override));
+  MOCK_METHOD(bool,
+              GetDmDevicePathByName,
+              (const std::string&, std::string*),
+              (override));
+  MOCK_METHOD(std::unique_ptr<::android::fs_mgr::MetadataBuilder>,
+              LoadMetadataBuilder,
+              (const std::string&, uint32_t, uint32_t),
+              (override));
+  MOCK_METHOD(bool,
+              StoreMetadata,
+              (const std::string&, android::fs_mgr::MetadataBuilder*, uint32_t),
+              (override));
+  MOCK_METHOD(bool, GetDeviceDir, (std::string*), (override));
+  MOCK_METHOD(FeatureFlag, GetDynamicPartitionsFeatureFlag, (), (override));
+  MOCK_METHOD(std::string, GetSuperPartitionName, (uint32_t), (override));
+  MOCK_METHOD(FeatureFlag, GetVirtualAbFeatureFlag, (), (override));
+  MOCK_METHOD(bool, FinishUpdate, (bool), (override));
+  MOCK_METHOD(bool,
+              GetSystemOtherPath,
+              (uint32_t, uint32_t, const std::string&, std::string*, bool*),
+              (override));
+  MOCK_METHOD(bool,
+              EraseSystemOtherAvbFooter,
+              (uint32_t, uint32_t),
+              (override));
+  MOCK_METHOD(std::optional<bool>, IsAvbEnabledOnSystemOther, (), (override));
+  MOCK_METHOD(bool, IsRecovery, (), (override));
+  MOCK_METHOD(bool,
+              PrepareDynamicPartitionsForUpdate,
+              (uint32_t, uint32_t, const DeltaArchiveManifest&, bool),
+              (override));
 
   void set_fake_mapped_devices(const std::set<std::string>& fake) override {
     DynamicPartitionControlAndroid::set_fake_mapped_devices(fake);
@@ -105,6 +100,15 @@
   std::optional<bool> RealIsAvbEnabledInFstab(const std::string& path) {
     return DynamicPartitionControlAndroid::IsAvbEnabledInFstab(path);
   }
+
+  bool RealPrepareDynamicPartitionsForUpdate(
+      uint32_t source_slot,
+      uint32_t target_slot,
+      const DeltaArchiveManifest& manifest,
+      bool delete_source) {
+    return DynamicPartitionControlAndroid::PrepareDynamicPartitionsForUpdate(
+        source_slot, target_slot, manifest, delete_source);
+  }
 };
 
 }  // namespace chromeos_update_engine
diff --git a/test_config.xml b/test_config.xml
index 2639e7f..fe3cbfd 100644
--- a/test_config.xml
+++ b/test_config.xml
@@ -16,13 +16,14 @@
 <configuration description="Config to run update_engine_unittests on device">
     <option name="test-suite-tag" value="apct" />
     <option name="test-suite-tag" value="apct-native" />
+    <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer" />
     <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
         <option name="cleanup" value="true" />
-        <option name="push" value="update_engine_unittests->/data/local/tmp/update_engine_unittests" />
+        <option name="push" value="update_engine_unittests->/data/nativetest/update_engine_unittests" />
     </target_preparer>
 
     <test class="com.android.tradefed.testtype.GTest" >
-        <option name="native-test-device-path" value="/data/local/tmp" />
+        <option name="native-test-device-path" value="/data/nativetest" />
         <!-- The following rules avoid test runner from calling the following helper executables
              directly as gtests. -->
         <option name="file-exclusion-filter-regex" value=".*/delta_generator$" />