[moblab] Use credentials file from the bucket if available. If the partner bucket has a credentials file - use that file and enable the cloud notifications. Fix issue with the command to restart after the wizard, in some cases restart was not a good option as the service had not started stop and start solves that issue. Cleanup, only validate the boto key one time, since it is now a more expensive operation. TEST=unit tests, manually tested running GTS on local moblab BUG=chromium:723863 Change-Id: I88e6d6901227fa667299a7c973d4f9276aceaf96 Reviewed-on: https://chromium-review.googlesource.com/508157 Commit-Ready: Keith Haddow <[email protected]> Tested-by: Keith Haddow <[email protected]> Reviewed-by: Keith Haddow <[email protected]> Reviewed-by: Michael Tang <[email protected]>

diff --git a/frontend/afe/moblab_rpc_interface.py b/frontend/afe/moblab_rpc_interface.py
index d83cafc..0fd155a 100644
--- a/frontend/afe/moblab_rpc_interface.py
+++ b/frontend/afe/moblab_rpc_interface.py

@@ -41,6 +41,7 @@
 _GS_SECRET_ACCESS_KEY = 'gs_secret_access_key'
 _RESULT_STORAGE_SERVER = 'results_storage_server'
 _USE_EXISTING_BOTO_FILE = 'use_existing_boto_file'
+_CLOUD_NOTIFICATION_ENABLED = 'cloud_notification_enabled'
 
 # Location where dhcp leases are stored.
 _DHCPD_LEASES = '/var/lib/dhcp/dhcpd.leases'
@@ -391,9 +392,6 @@
 
     @param cloud_storage_info: The JSON RPC object for cloud storage info.
     """
-    valid, details = _validate_cloud_storage_info(cloud_storage_info)
-    if not valid:
-        return _create_operation_status_response(valid, details)
     config_update = {}
     config_update['CROS'] = [
         (_IMAGE_STORAGE_SERVER, cloud_storage_info[_IMAGE_STORAGE_SERVER]),
@@ -411,13 +409,21 @@
         _write_config_file(MOBLAB_BOTO_LOCATION, boto_config, True)
 
     _CONFIG.parse_config_file()
+    _enable_notification_using_credentials_in_bucket()
     services = ['moblab-devserver-init',
-    'moblab-devserver-cleanup-init', ' moblab-gsoffloader_s-init',
-    'moblab-base-container-init', 'moblab-scheduler-init',
-    'moblab-gsoffloader-init']
-    cmd = ';sudo /sbin/restart '.join(services)
-    cmd += ';sudo /usr/sbin/apache2 -k graceful'
-    os.system("export ATEST_RESULTS_DIR=/usr/local/autotest/results;" + cmd)
+    'moblab-devserver-cleanup-init', 'moblab-gsoffloader_s-init',
+    'moblab-scheduler-init', 'moblab-gsoffloader-init']
+    cmd = 'export ATEST_RESULTS_DIR=/usr/local/autotest/results;'
+    cmd += 'sudo stop ' + ';sudo stop '.join(services)
+    cmd += ';sudo start ' + ';sudo start '.join(services)
+    cmd += ';sudo apache2 -k graceful'
+    logging.info(cmd)
+    try:
+        utils.run(cmd)
+    except error.CmdError as e:
+        logging.error(e)
+        # if all else fails reboot the device.
+        utils.run('sudo reboot')
 
     return _create_operation_status_response(True, None)
 
@@ -740,3 +746,23 @@
     pool=pool, run_prod_code=False, test_source_build=build,
     wait_for_results=False)
 
+
+def _enable_notification_using_credentials_in_bucket():
+    """ Check and enable cloud notification if a credentials file exits.
+    @return: None
+    """
+    gs_image_location =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
+    try:
+        utils.run(_GSUTIL_CMD, args=(
+            'cp', gs_image_location + 'pubsub-key-do-not-delete.json', '/tmp'))
+        # This runs the copy as moblab user
+        shutil.copyfile('/tmp/pubsub-key-do-not-delete.json',
+                        moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
+
+    except error.CmdError as e:
+        logging.error(e)
+    else:
+        logging.info('Enabling cloud notifications')
+        config_update = {}
+        config_update['CROS'] = [(_CLOUD_NOTIFICATION_ENABLED, True)]
+        _update_partial_config(config_update)

diff --git a/frontend/afe/moblab_rpc_interface_unittest.py b/frontend/afe/moblab_rpc_interface_unittest.py
index 062b061..eccdaac 100644
--- a/frontend/afe/moblab_rpc_interface_unittest.py
+++ b/frontend/afe/moblab_rpc_interface_unittest.py

@@ -440,6 +440,39 @@
             'key', 'secret', 'gs://bucket1', '1K', '10', 'testoutput')
         self.mox.VerifyAll()
 
+    def testEnableNotificationUsingCredentialsInBucketFail(self):
+        config_mock = self.mox.CreateMockAnything()
+        moblab_rpc_interface._CONFIG = config_mock
+        config_mock.get_config_value(
+            'CROS', 'image_storage_server').AndReturn('gs://bucket1/')
+        self.mox.StubOutWithMock(common_lib_utils, 'run')
+        common_lib_utils.run(moblab_rpc_interface._GSUTIL_CMD,
+            args=('cp', 'gs://bucket1/pubsub-key-do-not-delete.json',
+            '/tmp')).AndRaise(
+                error.CmdError("fakecommand", common_lib_utils.CmdResult(), ""))
+        self.mox.ReplayAll()
+        moblab_rpc_interface._enable_notification_using_credentials_in_bucket()
+
+    def testEnableNotificationUsingCredentialsInBucketSuccess(self):
+        config_mock = self.mox.CreateMockAnything()
+        moblab_rpc_interface._CONFIG = config_mock
+        config_mock.get_config_value(
+            'CROS', 'image_storage_server').AndReturn('gs://bucket1/')
+        self.mox.StubOutWithMock(common_lib_utils, 'run')
+        common_lib_utils.run(moblab_rpc_interface._GSUTIL_CMD,
+            args=('cp', 'gs://bucket1/pubsub-key-do-not-delete.json',
+            '/tmp'))
+        moblab_rpc_interface.shutil = self.mox.CreateMockAnything()
+        moblab_rpc_interface.shutil.copyfile(
+                '/tmp/pubsub-key-do-not-delete.json',
+                moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
+        self.mox.StubOutWithMock(moblab_rpc_interface, '_update_partial_config')
+        moblab_rpc_interface._update_partial_config(
+            {'CROS': [(moblab_rpc_interface._CLOUD_NOTIFICATION_ENABLED, True)]}
+        )
+        self.mox.ReplayAll()
+        moblab_rpc_interface._enable_notification_using_credentials_in_bucket()
+
 
 if __name__ == '__main__':
     unittest.main()