Check device build id and flavor after flashing

Bug: 17563593
Change-Id: I0b5027e6eab83a4f74c9f52618d966f411853ee3
diff --git a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
index 50d1842..b9b4ce8 100644
--- a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
+++ b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
@@ -121,7 +121,7 @@
         synchronized(sShouldCheckFlashLock) {
             // Make a minimal attempt to avoid having things get into an inconsistent state
             if (sConcurrentFlashLock != null && mConcurrentFlashLimit != null) {
-                int curLimit = (int) mConcurrentFlashLimit;
+                int curLimit = mConcurrentFlashLimit;
                 int curAvail = sConcurrentFlashLock.availablePermits();
                 if (curLimit != curAvail) {
                     throw new IllegalStateException(String.format("setConcurrentFlashSettings may " +
@@ -214,6 +214,7 @@
         if (device.enableAdbRoot()) {
             device.setDate(null);
         }
+        checkBuild(device, deviceBuild);
         postEncryptDevice(device, flasher);
         // only want logcat captured for current build, delete any accumulated log data
         device.clearLogcat();
@@ -230,6 +231,29 @@
     }
 
     /**
+     * Verifies the expected build matches the actual build on device after flashing
+     * @throws DeviceNotAvailableException
+     */
+    private void checkBuild(ITestDevice device, IDeviceBuildInfo deviceBuild)
+            throws DeviceNotAvailableException {
+        checkBuildAttribute(deviceBuild.getBuildId(), device.getBuildId());
+        checkBuildAttribute(deviceBuild.getBuildFlavor(), device.getBuildFlavor());
+        // TODO: check bootloader and baseband versions too
+    }
+
+    private void checkBuildAttribute(String expectedBuildAttr, String actualBuildAttr)
+            throws DeviceNotAvailableException {
+        if (expectedBuildAttr == null || actualBuildAttr == null ||
+                !expectedBuildAttr.equals(actualBuildAttr)) {
+            // throw DNAE - assume device hardware problem - we think flash was successful but
+            // device is not running right bits
+            throw new DeviceNotAvailableException(
+                    String.format("Unexpected build after flashing. Expected %s, actual %s",
+                            expectedBuildAttr, actualBuildAttr));
+        }
+    }
+
+    /**
      * Create {@link IDeviceFlasher} to use. Subclasses can override
      * @throws DeviceNotAvailableException
      */
diff --git a/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java b/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
index 45606df..4a112e6 100644
--- a/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
+++ b/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
@@ -23,7 +23,6 @@
 import com.android.tradefed.device.DeviceUnresponsiveException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.ITestDevice.RecoveryMode;
-import com.android.tradefed.targetprep.IDeviceFlasher;
 import com.android.tradefed.targetprep.IDeviceFlasher.UserDataFlashOption;
 import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
@@ -57,6 +56,7 @@
         mMockDevice = EasyMock.createMock(ITestDevice.class);
         EasyMock.expect(mMockDevice.getSerialNumber()).andReturn("foo").anyTimes();
         mMockBuildInfo = new DeviceBuildInfo("0", "", "");
+        mMockBuildInfo.setBuildFlavor("flavor");
         mDeviceFlashPreparer = new DeviceFlashPreparer() {
             @Override
             protected IDeviceFlasher createFlasher(ITestDevice device) {
@@ -105,6 +105,8 @@
         mMockDevice.waitForDeviceOnline();
         EasyMock.expect(mMockDevice.enableAdbRoot()).andStubReturn(Boolean.TRUE);
         mMockDevice.setDate(null);
+        EasyMock.expect(mMockDevice.getBuildId()).andReturn(mMockBuildInfo.getBuildId());
+        EasyMock.expect(mMockDevice.getBuildFlavor()).andReturn(mMockBuildInfo.getBuildFlavor());
         EasyMock.expect(mMockDevice.isEncryptionSupported()).andStubReturn(Boolean.TRUE);
         EasyMock.expect(mMockDevice.isDeviceEncrypted()).andStubReturn(Boolean.FALSE);
         mMockDevice.clearLogcat();
@@ -138,6 +140,8 @@
         mMockDevice.waitForDeviceOnline();
         EasyMock.expect(mMockDevice.enableAdbRoot()).andStubReturn(Boolean.TRUE);
         mMockDevice.setDate(null);
+        EasyMock.expect(mMockDevice.getBuildId()).andReturn(mMockBuildInfo.getBuildId());
+        EasyMock.expect(mMockDevice.getBuildFlavor()).andReturn(mMockBuildInfo.getBuildFlavor());
         EasyMock.expect(mMockDevice.isEncryptionSupported()).andStubReturn(Boolean.TRUE);
         EasyMock.expect(mMockDevice.isDeviceEncrypted()).andStubReturn(Boolean.FALSE);
         mMockDevice.clearLogcat();