[moblab] Add update button

BUG=chromium:613747
TEST=Tested features on a local moblab

Change-Id: Id512a96e57c5449ace561b8a551b2570ab637468
Reviewed-on: https://chromium-review.googlesource.com/887794
Commit-Ready: Matt Mallett <mattmallett@chromium.org>
Tested-by: Matt Mallett <mattmallett@chromium.org>
Reviewed-by: Keith Haddow <haddowk@chromium.org>
diff --git a/frontend/afe/moblab_rpc_interface.py b/frontend/afe/moblab_rpc_interface.py
index e368f4d..76fc227 100644
--- a/frontend/afe/moblab_rpc_interface.py
+++ b/frontend/afe/moblab_rpc_interface.py
@@ -53,6 +53,9 @@
 # File where information about the current device is stored.
 _ETC_LSB_RELEASE = '/etc/lsb-release'
 
+# ChromeOS update engine client binary location
+_UPDATE_ENGINE_CLIENT = '/usr/bin/update_engine_client'
+
 # Full path to the correct gsutil command to run.
 class GsUtil:
     """Helper class to find correct gsutil command."""
@@ -463,10 +466,73 @@
     version_response['MOBLAB_ID'] = utils.get_moblab_id();
     version_response['MOBLAB_MAC_ADDRESS'] = (
         utils.get_default_interface_mac_address())
+    _check_for_system_update()
+    update_status = _get_system_update_status()
+    version_response['MOBLAB_UPDATE_VERSION'] = update_status['NEW_VERSION']
+    version_response['MOBLAB_UPDATE_STATUS'] = update_status['CURRENT_OP']
+    version_response['MOBLAB_UPDATE_PROGRESS'] = update_status['PROGRESS']
     return rpc_utils.prepare_for_serialization(version_response)
 
 
 @rpc_utils.moblab_only
+def update_moblab():
+    """ RPC call to update and reboot moblab """
+    _install_system_update()
+
+
+def _check_for_system_update():
+    """ Run the ChromeOS update client to check update server for an
+    update. If an update exists, the update client begins downloading it
+    in the background
+    """
+    # sudo is required to run the update client
+    subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--check_for_update'])
+
+
+def _get_system_update_status():
+    """ Run the ChromeOS update client to check status on a
+    pending/downloading update
+
+    @return: A dictionary containing {
+        PROGRESS: str containing percent progress of an update download
+        CURRENT_OP: str current status of the update engine,
+            ex UPDATE_STATUS_UPDATED_NEED_REBOOT
+        NEW_SIZE: str size of the update
+        NEW_VERSION: str version number for the update
+        LAST_CHECKED_TIME: str unix time stamp of the last update check
+    }
+    """
+    # sudo is required to run the update client
+    cmd_out = subprocess.check_output(
+        ['sudo' ,_UPDATE_ENGINE_CLIENT, '--status'])
+    split_lines = [x.split('=') for x in cmd_out.strip().split('\n')]
+    status = dict((key, val) for [key, val] in split_lines)
+    return status
+
+
+def _install_system_update():
+    """ Installs a ChromeOS update, will cause the system to reboot
+    """
+    # sudo is required to run the update client
+    # first run a blocking command to check, fetch, prepare an update
+    # then check if a reboot is needed
+    try:
+        subprocess.check_call(['sudo', _UPDATE_ENGINE_CLIENT, '--update'])
+        try:
+            # --is_reboot_needed returns 2 if a reboot is required, which
+            # technically is an error
+            subprocess.check_call(
+                ['sudo', _UPDATE_ENGINE_CLIENT, '--is_reboot_needed'])
+        except subprocess.CalledProcessError as e:
+            if e.returncode == 2:
+                subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--reboot'])
+
+    except subprocess.CalledProcessError as e:
+        pass
+        #TODO(crbug/806311) surface error to UI
+
+
+@rpc_utils.moblab_only
 def get_connected_dut_info():
     """ RPC handler to get informaiton about the DUTs connected to the moblab.
 
diff --git a/frontend/afe/moblab_rpc_interface_unittest.py b/frontend/afe/moblab_rpc_interface_unittest.py
index 03b554d..b2bb296 100644
--- a/frontend/afe/moblab_rpc_interface_unittest.py
+++ b/frontend/afe/moblab_rpc_interface_unittest.py
@@ -490,6 +490,56 @@
         self.mox.ReplayAll()
         moblab_rpc_interface._enable_notification_using_credentials_in_bucket()
 
+    def testInstallSystemUpdate(self):
+        update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
+
+        self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'check_call')
+        moblab_rpc_interface.subprocess.check_call(['sudo',
+                update_engine_client, '--update'])
+        error = moblab_rpc_interface.subprocess.CalledProcessError(2, '')
+        moblab_rpc_interface.subprocess.check_call(['sudo',
+                update_engine_client, '--is_reboot_needed']).AndRaise(error)
+
+        self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call')
+        moblab_rpc_interface.subprocess.call(['sudo', update_engine_client,
+                '--reboot'])
+
+        self.mox.ReplayAll()
+        moblab_rpc_interface._install_system_update()
+
+
+    def testGetSystemUpdateStatus(self):
+        update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
+        update_status = ('LAST_CHECKED_TIME=1516753795\n'
+                         'PROGRESS=0.220121\n'
+                         'CURRENT_OP=UPDATE_STATUS_DOWNLOADING\n'
+                         'NEW_VERSION=10032.89.0\n'
+                         'NEW_SIZE=782805733')
+
+        self.mox.StubOutWithMock(moblab_rpc_interface.subprocess,
+                'check_output')
+        moblab_rpc_interface.subprocess.check_output(['sudo',
+                update_engine_client, '--status']).AndReturn(
+                        update_status)
+
+        self.mox.ReplayAll()
+        output = moblab_rpc_interface._get_system_update_status()
+
+        self.assertEquals(output['PROGRESS'], '0.220121')
+        self.assertEquals(output['CURRENT_OP'], 'UPDATE_STATUS_DOWNLOADING')
+        self.assertEquals(output['NEW_VERSION'], '10032.89.0')
+        self.assertEquals(output['NEW_SIZE'], '782805733')
+
+    def testCheckForSystemUpdate(self):
+        update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
+
+        self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call')
+        moblab_rpc_interface.subprocess.call(['sudo', update_engine_client,
+                '--check_for_update'])
+
+        self.mox.ReplayAll()
+        moblab_rpc_interface._check_for_system_update()
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/frontend/client/src/autotest/moblab/rpc/MoblabRpcHelper.java b/frontend/client/src/autotest/moblab/rpc/MoblabRpcHelper.java
index 0d59590..d279f1c 100644
--- a/frontend/client/src/autotest/moblab/rpc/MoblabRpcHelper.java
+++ b/frontend/client/src/autotest/moblab/rpc/MoblabRpcHelper.java
@@ -164,6 +164,14 @@
         });
   }
 
+  /**
+   * Apply update and reboot Moblab device
+   */
+  public static void updateMoblab(final JsonRpcCallback callback) {
+    JsonRpcProxy rpcProxy = JsonRpcProxy.getProxy();
+    rpcProxy.rpcCall("update_moblab", null, callback);
+  }
+
    /**
    * Get information about the DUT's connected to the moblab.
    */
diff --git a/frontend/client/src/autotest/moblab/rpc/VersionInfo.java b/frontend/client/src/autotest/moblab/rpc/VersionInfo.java
index b4f0c0a..0122cb6 100644
--- a/frontend/client/src/autotest/moblab/rpc/VersionInfo.java
+++ b/frontend/client/src/autotest/moblab/rpc/VersionInfo.java
@@ -7,12 +7,23 @@
 
 public class VersionInfo extends JsonRpcEntity {
 
+  public enum UPDATE_STATUS {
+    IDLE,
+    CHECKING_FOR_UPDATE,
+    UPDATE_AVAILABLE,
+    DOWNLOADING,
+    UPDATED_NEED_REBOOT,
+    UNKNOWN
+  }
+
   private static final String NO_MILESTONE_FOUND = "NO MILESTONE FOUND";
   private static final String NO_VERSION_FOUND = "NO VERSION FOUND";
   private static final String NO_TRACK_FOUND = "NO TRACK FOUND";
   private static final String NO_DESCRIPTION_FOUND = "NO DESCRIPTION FOUND";
   private static final String NO_ID_FOUND = "NO ID FOUND";
   private static final String NO_MAC_ADDRESS_FOUND = "NO MAC ADDRESS FOUND";
+  private static final String NO_UPDATE_VERSION_FOUND =
+      "NO UPDATE VERSION FOUND";
 
   private String milestoneInfo;
   private String versionInfo;
@@ -20,6 +31,9 @@
   private String releaseDescription;
   private String moblabIdentification;
   private String moblabMacAddress;
+  private String moblabUpdateVersion;
+  private double moblabUpdateProgress;
+  private UPDATE_STATUS moblabUpdateStatus;
 
   public VersionInfo() { reset(); }
 
@@ -30,6 +44,9 @@
   public String getReleaseDescription() { return releaseDescription; }
   public String getMoblabIdentification() { return moblabIdentification; }
   public String getMoblabMacAddress() { return moblabMacAddress; }
+  public String getMoblabUpdateVersion() { return moblabUpdateVersion; }
+  public double getMoblabUpdateProgress() { return moblabUpdateProgress; }
+  public UPDATE_STATUS getMoblabUpdateStatus() { return moblabUpdateStatus; }
 
   private void reset() {
     milestoneInfo = new String(NO_MILESTONE_FOUND);
@@ -38,6 +55,9 @@
     releaseDescription = new String(NO_DESCRIPTION_FOUND);
     moblabIdentification = new String(NO_ID_FOUND);
     moblabMacAddress = new String(NO_MAC_ADDRESS_FOUND);
+    moblabUpdateVersion = new String(NO_UPDATE_VERSION_FOUND);
+    moblabUpdateStatus = UPDATE_STATUS.UNKNOWN;
+    moblabUpdateProgress = 0.0;
   }
 
   @Override
@@ -55,6 +75,63 @@
         NO_DESCRIPTION_FOUND).trim();
     moblabMacAddress = getStringFieldOrDefault(object, "MOBLAB_MAC_ADDRESS",
         NO_DESCRIPTION_FOUND).trim();
+    moblabUpdateVersion = getStringFieldOrDefault(
+        object, "MOBLAB_UPDATE_VERSION", NO_UPDATE_VERSION_FOUND).trim();
+    moblabUpdateStatus = getUpdateStatus(object);
+
+    String progressString = getStringFieldOrDefault(
+        object, "MOBLAB_UPDATE_PROGRESS", "0.0").trim();
+    try {
+      moblabUpdateProgress = Double.parseDouble(progressString);
+    }
+    catch (NumberFormatException e) {
+      moblabUpdateProgress = 0.0;
+    }
+  }
+
+  private UPDATE_STATUS getUpdateStatus(JSONObject object) {
+    String status = getStringFieldOrDefault(
+        object, "MOBLAB_UPDATE_STATUS", "").trim();
+
+    if(status.contains("IDLE")) {
+      return UPDATE_STATUS.IDLE;
+    }
+    else if(status.contains("CHECKING_FOR_UPDATE")) {
+      return UPDATE_STATUS.CHECKING_FOR_UPDATE;
+    }
+    else if(status.contains("UPDATE_AVAILABLE")) {
+      return UPDATE_STATUS.UPDATE_AVAILABLE;
+    }
+    else if(status.contains("DOWNLOADING") || status.contains("VERIFYING") ||
+        status.contains("FINALIZING")) {
+      return UPDATE_STATUS.DOWNLOADING;
+    }
+    else if(status.contains("NEED_REBOOT")) {
+      return UPDATE_STATUS.UPDATED_NEED_REBOOT;
+    }
+    else {
+      return UPDATE_STATUS.UNKNOWN;
+    }
+  }
+
+  public String getUpdateString() {
+    switch(moblabUpdateStatus){
+      case CHECKING_FOR_UPDATE:
+        return "Checking for update..";
+      case UPDATE_AVAILABLE:
+        return "Version " + moblabUpdateVersion + " is available";
+      case DOWNLOADING:
+        int percent = (int)(moblabUpdateProgress * 100.0);
+        return "Downloading version " + moblabUpdateVersion
+            + " (" + percent + "%)";
+      case UPDATED_NEED_REBOOT:
+        return "Version " + moblabUpdateVersion
+            + " is available, reboot required";
+      case IDLE:
+      case UNKNOWN:
+      default:
+        return "";
+    }
   }
 
   @Override
@@ -64,4 +141,3 @@
     return new JSONObject();
   }
 }
-
diff --git a/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java b/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
index e24e531..ba6d619 100644
--- a/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
+++ b/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
@@ -11,6 +11,7 @@
 import com.google.gwt.user.client.ui.FlexTable;
 import com.google.gwt.user.client.ui.FlowPanel;
 import com.google.gwt.user.client.ui.Label;
+import com.google.gwt.user.client.ui.InlineLabel;
 import com.google.gwt.user.client.ui.SimplePanel;
 import com.google.gwt.user.client.ui.VerticalPanel;
 import com.google.gwt.user.client.ui.Widget;
@@ -122,6 +123,31 @@
         layoutTable.setWidget(row, 0, new Label("Version"));
         layoutTable.setWidget(row, 1, new Label(info.getVersion()));
         row++;
+
+        layoutTable.setWidget(row, 0, new Label("Update"));
+        FlowPanel updatePanel = new FlowPanel();
+        updatePanel.add(new InlineLabel(info.getUpdateString()));
+        Button btnUpdate = new Button("Update Now");
+        btnUpdate.addClickHandler(new ClickHandler() {
+          @Override
+          public void onClick(ClickEvent event) {
+            String windowText = "If an update is available, the device will be "
+              + "rebooted and all running jobs will be halted. Proceed?";
+            if (Window.confirm(windowText)) {
+              MoblabRpcHelper.updateMoblab(new JsonRpcCallback() {
+                @Override
+                public void onSuccess(JSONValue result) {
+                  String messageText = "Update command has been issued";
+                  NotifyManager.getInstance().showMessage(messageText);
+                }
+              });
+            }
+          }
+        });
+        updatePanel.add(btnUpdate);
+        layoutTable.setWidget(row, 1, updatePanel);
+        row++;
+
         layoutTable.setWidget(row, 0, new Label("Track"));
         layoutTable.setWidget(row, 1, new Label(info.getReleaseTrack()));
         row++;