Enable OTA build info to report its target build

For some OTA tests, the instance of OtaDeviceBuildInfo needs to
report its 'primary' build as the target build rather than the
source build.

Bug: 31777654
Change-Id: I6ed7f93e558fe740d85b2bbe1c91e6064c40e177

2 files changed, 25 insertions, 2 deletions

diff --git a/src/com/android/tradefed/build/OtaDeviceBuildInfo.java b/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
index 0011c3250..71c73a432 100644
--- a/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
@@ -36,6 +36,7 @@ public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
 
     protected IDeviceBuildInfo mOtaBuild;
     protected IDeviceBuildInfo mBaselineBuild;
+    protected boolean mDowngrade = false;
 
     public void setOtaBuild(IDeviceBuildInfo otaBuild) {
         mOtaBuild = otaBuild;
@@ -54,6 +55,9 @@ public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
      */
     @Override
     public String getBuildId() {
+        if (mDowngrade) {
+            return mOtaBuild.getBuildId();
+        }
         return mBaselineBuild.getBuildId();
     }
 
@@ -198,6 +202,9 @@ public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
      */
     @Override
     public File getDeviceImageFile() {
+        if (mDowngrade) {
+            return mOtaBuild.getDeviceImageFile();
+        }
         return mBaselineBuild.getDeviceImageFile();
     }
 
@@ -206,6 +213,9 @@ public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
      */
     @Override
     public String getDeviceImageVersion() {
+        if (mDowngrade) {
+            return mOtaBuild.getDeviceImageVersion();
+        }
         return mBaselineBuild.getDeviceImageVersion();
     }
 
@@ -419,4 +429,8 @@ public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
         combinedFiles.addAll(mOtaBuild.getFiles());
         return combinedFiles;
     }
+
+    public void setDowngrade(boolean downgrade) {
+        mDowngrade = downgrade;
+    }
 }
diff --git a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
index c0276b7a2..ed265017d 100644
--- a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
+++ b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
@@ -70,6 +70,13 @@ public abstract class DeviceFlashPreparer implements ITargetCleaner {
             "specify if system flavor should not be checked after flash")
     private boolean mSkipPostFlashFlavorCheck = false;
 
+    /*
+     * Used for update testing
+     */
+    @Option(name = "skip-post-flash-build-id-check", description =
+            "specify if build ID should not be checked after flash")
+    private boolean mSkipPostFlashBuildIdCheck = false;
+
     @Option(name = "wipe-skip-list", description =
         "list of /data subdirectories to NOT wipe when doing UserDataFlashOption.TESTS_ZIP")
     private Collection<String> mDataWipeSkipList = new ArrayList<>();
@@ -304,8 +311,10 @@ public abstract class DeviceFlashPreparer implements ITargetCleaner {
         // Need to use deviceBuild.getDeviceBuildId instead of getBuildId because the build info
         // could be an AppBuildInfo and return app build id. Need to be more explicit that we
         // check for the device build here.
-        checkBuildAttribute(deviceBuild.getDeviceBuildId(), device.getBuildId(),
-                device.getSerialNumber());
+        if (!mSkipPostFlashBuildIdCheck) {
+            checkBuildAttribute(deviceBuild.getDeviceBuildId(), device.getBuildId(),
+                    device.getSerialNumber());
+        }
         if (!mSkipPostFlashFlavorCheck) {
             checkBuildAttribute(deviceBuild.getDeviceBuildFlavor(), device.getBuildFlavor(),
                     device.getSerialNumber());