Remove boilerplate setup code around nebraska_wrapper.

Change args in nebraksa_wrapper and update the callers.
Now just give it the payload_url and the props to override and let it do
everything instead of each test doing it.

BUG=chromium:1081031
TEST=backoff,badmetadata,badsignatures,nonblockingoobe,ui,omaharesponse.

Change-Id: If72281516eadbdcc56e8e6a3230fe24d1d9a84dc
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2224911
Tested-by: David Haddock <dhaddock@chromium.org>
Reviewed-by: Amin Hassani <ahassani@chromium.org>
Commit-Queue: David Haddock <dhaddock@chromium.org>
diff --git a/client/cros/update_engine/nebraska_wrapper.py b/client/cros/update_engine/nebraska_wrapper.py
index 498ccc5..89f7c77 100644
--- a/client/cros/update_engine/nebraska_wrapper.py
+++ b/client/cros/update_engine/nebraska_wrapper.py
@@ -2,6 +2,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import json
 import logging
 import os
 import requests
@@ -9,6 +10,7 @@
 import urlparse
 
 from autotest_lib.client.bin import utils
+from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 
 
@@ -28,22 +30,29 @@
 
     """
 
-    def __init__(self, log_dir=None, update_metadata_dir=None,
-                 update_payloads_address=None):
+    def __init__(self, log_dir=None, payload_url=None, **props_to_override):
         """
         Initializes the NebraskaWrapper module.
 
         @param log_dir: The directory to write nebraska.log into.
-        @param update_metadata_dir: The directory containing payload properties
-                files. Look at nebraska.py
-        @param update_payloads_address: The base URL for the update payload.
+        @param payload_url: The payload that will returned in responses.
+        @param props_to_override: Dictionary of key/values to use in responses
+                instead of the default values in payload_url's properties file.
 
         """
         self._nebraska_server = None
         self._port = None
         self._log_dir = log_dir
-        self._update_metadata_dir = update_metadata_dir
-        self._update_payloads_address = update_payloads_address
+        self._update_metadata_dir = None
+        self._update_payloads_address = None
+
+        if payload_url:
+            self._update_metadata_dir = autotemp.tempdir()
+            self._update_payloads_address = ''.join(
+                payload_url.rpartition('/')[0:2])
+            self.get_payload_properties_file(
+                payload_url, self._update_metadata_dir.name,
+                **props_to_override)
 
     def __enter__(self):
         """So that NebraskaWrapper can be used as a Context Manager."""
@@ -74,7 +83,7 @@
         if self._log_dir:
             cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')]
         if self._update_metadata_dir:
-            cmd += ['--update-metadata', self._update_metadata_dir]
+            cmd += ['--update-metadata', self._update_metadata_dir.name]
         if self._update_payloads_address:
             cmd += ['--update-payloads-address', self._update_payloads_address]
 
@@ -129,3 +138,33 @@
                                    query=query,
                                    fragment='')
         return urlparse.urlunsplit(url)
+
+    def get_payload_properties_file(self, payload_url, target_dir, **kwargs):
+        """
+        Downloads the payload properties file into a directory.
+
+        @param payload_url: The URL to the update payload file.
+        @param target_dir: The directory to download the file into.
+        @param kwargs: A dictionary of key/values that needs to be overridden on
+                the payload properties file.
+
+        """
+        payload_props_url = payload_url + '.json'
+        _, _, file_name = payload_props_url.rpartition('/')
+        try:
+            response = json.loads(requests.get(payload_props_url).text)
+            # Override existing keys if any.
+            for k, v in kwargs.iteritems():
+                # Don't set default None values. We don't want to override good
+                # values to None.
+                if v is not None:
+                    response[k] = v
+            with open(os.path.join(target_dir, file_name), 'w') as fp:
+                json.dump(response, fp)
+
+        except (requests.exceptions.RequestException,
+                IOError,
+                ValueError) as err:
+            raise error.TestError(
+                'Failed to get update payload properties: %s with error: %s' %
+                (payload_props_url, err))
diff --git a/client/cros/update_engine/update_engine_util.py b/client/cros/update_engine/update_engine_util.py
index 84a09d9..cd80fad 100644
--- a/client/cros/update_engine/update_engine_util.py
+++ b/client/cros/update_engine/update_engine_util.py
@@ -3,11 +3,9 @@
 # found in the LICENSE file.
 
 import datetime
-import json
 import logging
 import os
 import re
-import requests
 import shutil
 import time
 import urlparse
@@ -381,39 +379,6 @@
         return True
 
 
-    def _get_payload_properties_file(self, payload_url, target_dir, **kwargs):
-        """
-        Downloads the payload properties file into a directory.
-
-        @param payload_url: The URL to the update payload file.
-        @param target_dir: The directory to download the file into.
-        @param kwargs: A dictionary of key/values that needs to be overridden on
-                the payload properties file.
-
-        """
-        payload_props_url = payload_url + '.json'
-        _, _, file_name = payload_props_url.rpartition('/')
-        try:
-            response = json.loads(requests.get(payload_props_url).text)
-
-            # Override existing keys if any.
-            for k, v in kwargs.iteritems():
-                # Don't set default None values. We don't want to override good
-                # values to None.
-                if v is not None:
-                    response[k] = v
-
-            with open(os.path.join(target_dir, file_name), 'w') as fp:
-                json.dump(response, fp)
-
-        except (requests.exceptions.RequestException,
-                IOError,
-                ValueError) as err:
-            raise error.TestError(
-                'Failed to get update payload properties: %s with error: %s' %
-                (payload_props_url, err))
-
-
     def _append_query_to_url(self, url, query_dict):
         """
         Appends the dictionary kwargs to the URL url as query string.
diff --git a/client/site_tests/autoupdate_Backoff/autoupdate_Backoff.py b/client/site_tests/autoupdate_Backoff/autoupdate_Backoff.py
index 7b416ea..ff05b9a 100644
--- a/client/site_tests/autoupdate_Backoff/autoupdate_Backoff.py
+++ b/client/site_tests/autoupdate_Backoff/autoupdate_Backoff.py
@@ -6,7 +6,6 @@
 import os
 
 from autotest_lib.client.bin import utils
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
 from autotest_lib.client.cros.update_engine import update_engine_test
@@ -55,15 +54,8 @@
                                 self._NO_IGNORE_BACKOFF_PREF)],
                   ignore_status=True)
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url,
-                                          metadata_dir.name)
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
-
+            log_dir=self.resultsdir, payload_url=payload_url) as nebraska:
             # Only set one URL in the Nebraska response so we can test the
             # backoff functionality quicker.
             response_props = {'disable_payload_backoff': not backoff,
diff --git a/client/site_tests/autoupdate_BadMetadata/autoupdate_BadMetadata.py b/client/site_tests/autoupdate_BadMetadata/autoupdate_BadMetadata.py
index 6cb9395..35a08f6 100644
--- a/client/site_tests/autoupdate_BadMetadata/autoupdate_BadMetadata.py
+++ b/client/site_tests/autoupdate_BadMetadata/autoupdate_BadMetadata.py
@@ -4,7 +4,6 @@
 
 import logging
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
 from autotest_lib.client.cros.update_engine import update_engine_test
@@ -39,15 +38,9 @@
                 nebraska_wrapper.KEY_PUBLIC_KEY] = self._IMAGE_PUBLIC_KEY
             error_string = self._METADATA_SIZE_ERROR
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url,
-                                          metadata_dir.name,
-                                          **props_to_override)
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=payload_url,
+            **props_to_override) as nebraska:
 
             try:
                 self._check_for_update(
diff --git a/client/site_tests/autoupdate_CannedOmahaUpdate/autoupdate_CannedOmahaUpdate.py b/client/site_tests/autoupdate_CannedOmahaUpdate/autoupdate_CannedOmahaUpdate.py
index 59202e7..2175bc9 100644
--- a/client/site_tests/autoupdate_CannedOmahaUpdate/autoupdate_CannedOmahaUpdate.py
+++ b/client/site_tests/autoupdate_CannedOmahaUpdate/autoupdate_CannedOmahaUpdate.py
@@ -4,7 +4,6 @@
 
 import logging
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.cros.cellular import test_environment
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
@@ -48,15 +47,9 @@
 
         """
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url,
-                                          metadata_dir.name,
-                                          public_key=public_key)
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=payload_url,
+            public_key=public_key) as nebraska:
 
             if not use_cellular:
                 self.run_canned_update(allow_failure, nebraska.get_update_url())
diff --git a/client/site_tests/autoupdate_StartOOBEUpdate/autoupdate_StartOOBEUpdate.py b/client/site_tests/autoupdate_StartOOBEUpdate/autoupdate_StartOOBEUpdate.py
index 2dca3a9..72710cc 100644
--- a/client/site_tests/autoupdate_StartOOBEUpdate/autoupdate_StartOOBEUpdate.py
+++ b/client/site_tests/autoupdate_StartOOBEUpdate/autoupdate_StartOOBEUpdate.py
@@ -5,7 +5,6 @@
 import logging
 
 from autotest_lib.client.bin import utils
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import chrome
 from autotest_lib.client.cros.cellular import test_environment
@@ -113,16 +112,11 @@
                                          'it left off after interruption.')
             return
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url, metadata_dir.name)
-        # Setup a Nebraska instance on the DUT because we can't reach devservers
-        # over cellular. Same for non-critical updates which allow better
-        # debugging.
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
+        # Setup a Nebraska instance on the DUT for cellular tests and
+        # non-critical updates. Ceullar tests cannot reach devservers.
+        # Non-critical tests don't need a devserver.
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=payload_url) as nebraska:
 
             update_url = nebraska.get_update_url(
                 critical_update=critical_update)
diff --git a/client/site_tests/autoupdate_UpdateFromUI/autoupdate_UpdateFromUI.py b/client/site_tests/autoupdate_UpdateFromUI/autoupdate_UpdateFromUI.py
index 4441d9b..5a80308 100644
--- a/client/site_tests/autoupdate_UpdateFromUI/autoupdate_UpdateFromUI.py
+++ b/client/site_tests/autoupdate_UpdateFromUI/autoupdate_UpdateFromUI.py
@@ -2,7 +2,6 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import chrome
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
@@ -33,13 +32,8 @@
         @param payload_url: The payload url to use.
 
         """
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url, metadata_dir.name)
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=payload_url) as nebraska:
             with chrome.Chrome(autotest_ext=True) as cr:
                 # Need to create a custom lsb-release file to point the UI
                 # update button to Nebraska instead of the default update
diff --git a/client/site_tests/autoupdate_UrlSwitch/autoupdate_UrlSwitch.py b/client/site_tests/autoupdate_UrlSwitch/autoupdate_UrlSwitch.py
index 2421e95..945a399 100644
--- a/client/site_tests/autoupdate_UrlSwitch/autoupdate_UrlSwitch.py
+++ b/client/site_tests/autoupdate_UrlSwitch/autoupdate_UrlSwitch.py
@@ -2,7 +2,6 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
 from autotest_lib.client.cros.update_engine import update_engine_test
 
@@ -18,14 +17,8 @@
         @param image_url: The URL of the update payload.
 
         """
-        # Get payload properties file so we can run Nebraska with it.
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(payload_url, metadata_dir.name)
-        base_url = ''.join(payload_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=payload_url) as nebraska:
 
             # Start the update that will return two Urls. This matches what test
             # and production omaha does today.
diff --git a/client/site_tests/autoupdate_UserData/autoupdate_UserData.py b/client/site_tests/autoupdate_UserData/autoupdate_UserData.py
index 40b55a4..e4ff51b 100644
--- a/client/site_tests/autoupdate_UserData/autoupdate_UserData.py
+++ b/client/site_tests/autoupdate_UserData/autoupdate_UserData.py
@@ -6,7 +6,6 @@
 import os
 
 from autotest_lib.client.bin import utils
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import chrome
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
@@ -111,13 +110,8 @@
 
         """
         if payload_url:
-            metadata_dir = autotemp.tempdir()
-            self._get_payload_properties_file(payload_url, metadata_dir.name)
-            base_url = ''.join(payload_url.rpartition('/')[0:2])
             with nebraska_wrapper.NebraskaWrapper(
-                    log_dir=self.resultsdir,
-                    update_metadata_dir=metadata_dir.name,
-                    update_payloads_address=base_url) as nebraska:
+                log_dir=self.resultsdir, payload_url=payload_url) as nebraska:
                 with chrome.Chrome(autotest_ext=True) as cr:
                     self._cr = cr
                     utils.run(['echo', 'hello', '>', self._TEST_FILE])
diff --git a/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py b/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py
index 308d43d..088c043 100644
--- a/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py
+++ b/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py
@@ -6,7 +6,6 @@
 import math
 import time
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib import utils
 from autotest_lib.client.cros.update_engine import nebraska_wrapper
@@ -77,14 +76,9 @@
 
         self.setup_case(device_policies={self._POLICY: case}, enroll=enroll)
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(image_url, metadata_dir.name,
-                                          target_version='999999.9.9')
-        base_url = ''.join(image_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=image_url,
+            target_version='999999.9.9') as nebraska:
 
             self._create_custom_lsb_release(nebraska.get_update_url(),
                                             build='1.1.1')
diff --git a/client/site_tests/policy_DeviceTargetVersionPrefix/policy_DeviceTargetVersionPrefix.py b/client/site_tests/policy_DeviceTargetVersionPrefix/policy_DeviceTargetVersionPrefix.py
index 7552b44..75e6ddd 100644
--- a/client/site_tests/policy_DeviceTargetVersionPrefix/policy_DeviceTargetVersionPrefix.py
+++ b/client/site_tests/policy_DeviceTargetVersionPrefix/policy_DeviceTargetVersionPrefix.py
@@ -5,7 +5,6 @@
 import logging
 import re
 
-from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib import utils
 from autotest_lib.client.cros.enterprise import enterprise_policy_base
@@ -98,14 +97,9 @@
         self.setup_case(device_policies={self._POLICY_NAME: case_value},
                         enroll=True)
 
-        metadata_dir = autotemp.tempdir()
-        self._get_payload_properties_file(image_url, metadata_dir.name,
-                                          target_version='999999.9.9')
-        base_url = ''.join(image_url.rpartition('/')[0:2])
         with nebraska_wrapper.NebraskaWrapper(
-                log_dir=self.resultsdir,
-                update_metadata_dir=metadata_dir.name,
-                update_payloads_address=base_url) as nebraska:
+            log_dir=self.resultsdir, payload_url=image_url,
+            target_version='999999.9.9') as nebraska:
 
             update_url = nebraska.get_update_url()
             self._create_custom_lsb_release(update_url, build='1.1.1')