Repository: Azure/WALinuxAgent
Pull request #399 (2016-08-30 22:55:39)

Use proxy for storage calls
- ensure all calls to `call_storage_service` use the http proxy
- pep8 fixes and cleanup
- addresses #364 

/cc @brendandixon 

brendandixon: @hglkrijger I added a line comment. Could we forcibly add the chi_proxy flag within the call_storage_service method?
hglkrijger: @brendandixon, sure. Initially that route seemed a bit heavy-handed to me, but after thinking about it some more I think it's the right call. Updated the PR.  
jasonzio: Now that chk_proxy is forced on, you don't need to supply it in all the calls.
hglkrijger: @jasonzio, updated the PR.
jasonzio: lgtm
hglkrijger: thanks

Patch:
diff --git a/azurelinuxagent/common/protocol/restapi.py b/azurelinuxagent/common/protocol/restapi.py
index 7f004884..3036925b 100644
--- a/azurelinuxagent/common/protocol/restapi.py
+++ b/azurelinuxagent/common/protocol/restapi.py
@@ -16,15 +16,12 @@
 #
 # Requires Python 2.4+ and Openssl 1.0+
 #
-import os
-import copy
-import re
-import json
-import xml.dom.minidom
+
 import azurelinuxagent.common.logger as logger
+import azurelinuxagent.common.utils.restutil as restutil
 from azurelinuxagent.common.exception import ProtocolError, HttpError
 from azurelinuxagent.common.future import ustr
-import azurelinuxagent.common.utils.restutil as restutil
+
 
 def validate_param(name, val, expected_type):
     if val is None:
@@ -33,13 +30,14 @@ def validate_param(name, val, expected_type):
         raise ProtocolError(("{0} type should be {1} not {2}"
                              "").format(name, expected_type, type(val)))
 
+
 def set_properties(name, obj, data):
     if isinstance(obj, DataContract):
         validate_param("Property '{0}'".format(name), data, dict)
         for prob_name, prob_val in data.items():
             prob_full_name = "{0}.{1}".format(name, prob_name)
             try:
-                prob =  getattr(obj, prob_name)
+                prob = getattr(obj, prob_name)
             except AttributeError:
                 logger.warn("Unknown property: {0}", prob_full_name)
                 continue
@@ -56,6 +54,7 @@ def set_properties(name, obj, data):
     else:
         return data
 
+
 def get_properties(obj):
     if isinstance(obj, DataContract):
         data = {}
@@ -72,16 +71,21 @@ def get_properties(obj):
     else:
         return obj
 
+
 class DataContract(object):
     pass
 
+
 class DataContractList(list):
     def __init__(self, item_cls):
         self.item_cls = item_cls
 
+
 """
 Data contract between guest and host
 """
+
+
 class VMInfo(DataContract):
     def __init__(self, subscriptionId=None, vmName=None, containerId=None,
                  roleName=None, roleInstanceName=None, tenantName=None):
@@ -92,30 +96,36 @@ class VMInfo(DataContract):
         self.roleInstanceName = roleInstanceName
         self.tenantName = tenantName
 
+
 class Cert(DataContract):
     def __init__(self, name=None, thumbprint=None, certificateDataUri=None):
         self.name = name
         self.thumbprint = thumbprint
         self.certificateDataUri = certificateDataUri
 
+
 class CertList(DataContract):
     def __init__(self):
         self.certificates = DataContractList(Cert)
 
-#TODO: confirm vmagent manifest schema
+
+# TODO: confirm vmagent manifest schema
 class VMAgentManifestUri(DataContract):
     def __init__(self, uri=None):
         self.uri = uri
 
+
 class VMAgentManifest(DataContract):
     def __init__(self, family=None):
         self.family = family
         self.versionsManifestUris = DataContractList(VMAgentManifestUri)
 
+
 class VMAgentManifestList(DataContract):
     def __init__(self):
         self.vmAgentManifests = DataContractList(VMAgentManifest)
 
+
 class Extension(DataContract):
     def __init__(self, name=None, sequenceNumber=None, publicSettings=None,
                  protectedSettings=None, certificateThumbprint=None):
@@ -125,6 +135,7 @@ class Extension(DataContract):
         self.protectedSettings = protectedSettings
         self.certificateThumbprint = certificateThumbprint
 
+
 class ExtHandlerProperties(DataContract):
     def __init__(self):
         self.version = None
@@ -132,40 +143,48 @@ class ExtHandlerProperties(DataContract):
         self.state = None
         self.extensions = DataContractList(Extension)
 
+
 class ExtHandlerVersionUri(DataContract):
     def __init__(self):
         self.uri = None
 
+
 class ExtHandler(DataContract):
     def __init__(self, name=None):
         self.name = name
         self.properties = ExtHandlerProperties()
         self.versionUris = DataContractList(ExtHandlerVersionUri)
 
+
 class ExtHandlerList(DataContract):
     def __init__(self):
         self.extHandlers = DataContractList(ExtHandler)
 
+
 class ExtHandlerPackageUri(DataContract):
     def __init__(self, uri=None):
         self.uri = uri
 
+
 class ExtHandlerPackage(DataContract):
-    def __init__(self, version = None):
+    def __init__(self, version=None):
         self.version = version
         self.uris = DataContractList(ExtHandlerPackageUri)
         # TODO update the naming to align with metadata protocol
         self.isinternal = False
 
+
 class ExtHandlerPackageList(DataContract):
     def __init__(self):
         self.versions = DataContractList(ExtHandlerPackage)
 
+
 class VMProperties(DataContract):
     def __init__(self, certificateThumbprint=None):
-        #TODO need to confirm the property name
+        # TODO need to confirm the property name
         self.certificateThumbprint = certificateThumbprint
 
+
 class ProvisionStatus(DataContract):
     def __init__(self, status=None, subStatus=None, description=None):
         self.status = status
@@ -173,6 +192,7 @@ class ProvisionStatus(DataContract):
         self.description = description
         self.properties = VMProperties()
 
+
 class ExtensionSubStatus(DataContract):
     def __init__(self, name=None, status=None, code=None, message=None):
         self.name = name
@@ -180,6 +200,7 @@ class ExtensionSubStatus(DataContract):
         self.code = code
         self.message = message
 
+
 class ExtensionStatus(DataContract):
     def __init__(self, configurationAppliedTime=None, operation=None,
                  status=None, seq_no=None, code=None, message=None):
@@ -191,8 +212,9 @@ class ExtensionStatus(DataContract):
         self.message = message
         self.substatusList = DataContractList(ExtensionSubStatus)
 
+
 class ExtHandlerStatus(DataContract):
-    def __init__(self, name=None, version=None, status=None, code=0, 
+    def __init__(self, name=None, version=None, status=None, code=0,
                  message=None):
         self.name = name
         self.version = version
@@ -201,6 +223,7 @@ class ExtHandlerStatus(DataContract):
         self.message = message
         self.extensions = DataContractList(ustr)
 
+
 class VMAgentStatus(DataContract):
     def __init__(self, version=None, status=None, message=None):
         self.version = version
@@ -208,27 +231,31 @@ class VMAgentStatus(DataContract):
         self.message = message
         self.extensionHandlers = DataContractList(ExtHandlerStatus)
 
+
 class VMStatus(DataContract):
     def __init__(self):
         self.vmAgent = VMAgentStatus()
 
+
 class TelemetryEventParam(DataContract):
     def __init__(self, name=None, value=None):
         self.name = name
         self.value = value
 
+
 class TelemetryEvent(DataContract):
     def __init__(self, eventId=None, providerId=None):
         self.eventId = eventId
         self.providerId = providerId
         self.parameters = DataContractList(TelemetryEventParam)
 
+
 class TelemetryEventList(DataContract):
     def __init__(self):
         self.events = DataContractList(TelemetryEvent)
 
-class Protocol(DataContract):
 
+class Protocol(DataContract):
     def detect(self):
         raise NotImplementedError()
 
@@ -240,7 +267,7 @@ class Protocol(DataContract):
 
     def get_vmagent_manifests(self):
         raise NotImplementedError()
-    
+
     def get_vmagent_pkgs(self):
         raise NotImplementedError()
 
@@ -269,4 +296,3 @@ class Protocol(DataContract):
 
     def report_event(self, event):
         raise NotImplementedError()
-
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index 287da33c..45f79f32 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -16,13 +16,16 @@
 #
 # Requires Python 2.4+ and Openssl 1.0+
 
+import json
+import os
+import re
 import time
 import xml.sax.saxutils as saxutils
 import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import ProtocolNotFoundError
 from azurelinuxagent.common.future import httpclient, bytebuffer
-from azurelinuxagent.common.utils.textutil import parse_doc, findall, find, findtext, \
-    getattrib, gettext, remove_bom, get_bytes_from_pem
+from azurelinuxagent.common.utils.textutil import parse_doc, findall, find, \
+    findtext, getattrib, gettext, remove_bom, get_bytes_from_pem
 import azurelinuxagent.common.utils.fileutil as fileutil
 from azurelinuxagent.common.utils.cryptutil import CryptUtil
 from azurelinuxagent.common.protocol.restapi import *
@@ -66,10 +69,13 @@ class WireProtocol(Protocol):
     """Slim layer to adapt wire protocol data to metadata protocol interface"""
 
     # TODO: Clean-up goal state processing
-    #   At present, some methods magically update GoalState (e.g., get_vmagent_manifests), others (e.g., get_vmagent_pkgs)
-    #   assume its presence. A better approach would make an explicit update call that returns the incarnation number and
-    #   establishes that number the "context" for all other calls (either by updating the internal state of the protocol or
-    #   by having callers pass the incarnation number to the method).
+    #  At present, some methods magically update GoalState (e.g.,
+    #  get_vmagent_manifests), others (e.g., get_vmagent_pkgs)
+    #  assume its presence. A better approach would make an explicit update
+    #  call that returns the incarnation number and
+    #  establishes that number the "context" for all other calls (either by
+    #  updating the internal state of the protocol or
+    #  by having callers pass the incarnation number to the method).
 
     def __init__(self, endpoint):
         if endpoint is None:
@@ -369,10 +375,13 @@ class StatusBlob(object):
         logger.verbose("Check blob type.")
         timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
         try:
-            resp = self.client.call_storage_service(restutil.http_head, url, {
-                "x-ms-date": timestamp,
-                'x-ms-version': self.__class__.__storage_version__
-            })
+            resp = self.client.call_storage_service(
+                restutil.http_head,
+                url,
+                {
+                    "x-ms-date": timestamp,
+                    'x-ms-version': self.__class__.__storage_version__
+                })
         except HttpError as e:
             raise ProtocolError((u"Failed to get status blob type: {0}"
                                  u"").format(e))
@@ -387,13 +396,16 @@ class StatusBlob(object):
     def put_block_blob(self, url, data):
         logger.verbose("Upload block blob")
         timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-        resp = self.client.call_storage_service(restutil.http_put, url, data,
-                        {
-                            "x-ms-date": timestamp,
-                            "x-ms-blob-type": "BlockBlob",
-                            "Content-Length": ustr(len(data)),
-                            "x-ms-version": self.__class__.__storage_version__
-                        })
+        resp = self.client.call_storage_service(
+            restutil.http_put,
+            url,
+            data,
+            {
+                "x-ms-date": timestamp,
+                "x-ms-blob-type": "BlockBlob",
+                "Content-Length": ustr(len(data)),
+                "x-ms-version": self.__class__.__storage_version__
+            })
         if resp.status != httpclient.CREATED:
             raise UploadError(
                 "Failed to upload block blob: {0}".format(resp.status))
@@ -407,14 +419,17 @@ class StatusBlob(object):
 
         # Align to 512 bytes
         page_blob_size = int((len(data) + 511) / 512) * 512
-        resp = self.client.call_storage_service(restutil.http_put, url, "",
-                        {
-                            "x-ms-date": timestamp,
-                            "x-ms-blob-type": "PageBlob",
-                            "Content-Length": "0",
-                            "x-ms-blob-content-length": ustr(page_blob_size),
-                            "x-ms-version": self.__class__.__storage_version__
-                        })
+        resp = self.client.call_storage_service(
+            restutil.http_put,
+            url,
+            "",
+            {
+                "x-ms-date": timestamp,
+                "x-ms-blob-type": "PageBlob",
+                "Content-Length": "0",
+                "x-ms-blob-content-length": ustr(page_blob_size),
+                "x-ms-version": self.__class__.__storage_version__
+            })
         if resp.status != httpclient.CREATED:
             raise UploadError(
                 "Failed to clean up page blob: {0}".format(resp.status))
@@ -437,7 +452,9 @@ class StatusBlob(object):
             buf = bytearray(buf_size)
             buf[0: content_size] = data[start: end]
             resp = self.client.call_storage_service(
-                restutil.http_put, url, bytebuffer(buf),
+                restutil.http_put,
+                url,
+                bytebuffer(buf),
                 {
                     "x-ms-date": timestamp,
                     "x-ms-range": "bytes={0}-{1}".format(start, page_end - 1),
@@ -465,7 +482,8 @@ def event_param_to_v1(param):
         attr_type = 'mt:bool'
     elif param_type is float:
         attr_type = 'mt:float64'
-    return param_format.format(param.name, saxutils.quoteattr(ustr(param.value)),
+    return param_format.format(param.name,
+                               saxutils.quoteattr(ustr(param.value)),
                                attr_type)
 
 
@@ -501,15 +519,15 @@ class WireClient(object):
         now = time.time()
         if now - self.last_request < 1:
             logger.verbose("Last request issued less than 1 second ago")
-            logger.verbose("Sleep {0} second to avoid throttling.", 
-                        SHORT_WAITING_INTERVAL)
+            logger.verbose("Sleep {0} second to avoid throttling.",
+                           SHORT_WAITING_INTERVAL)
             time.sleep(SHORT_WAITING_INTERVAL)
         self.last_request = now
 
         self.req_count += 1
         if self.req_count % 3 == 0:
-            logger.verbose("Sleep {0} second to avoid throttling.", 
-                        SHORT_WAITING_INTERVAL)
+            logger.verbose("Sleep {0} second to avoid throttling.",
+                           SHORT_WAITING_INTERVAL)
             time.sleep(SHORT_WAITING_INTERVAL)
             self.req_count = 0
 
@@ -566,29 +584,36 @@ class WireClient(object):
         except IOError as e:
             raise ProtocolError("Failed to write cache: {0}".format(e))
 
-    def call_storage_service(self, http_req, *args, **kwargs):
+    @staticmethod
+    def call_storage_service(http_req, *args, **kwargs):
         """ 
         Call storage service, handle SERVICE_UNAVAILABLE(503)
         """
+
+        # force the chk_proxy arg to True, since all calls to storage should
+        #  use a configured proxy
+        kwargs['chk_proxy'] = True
+
         for retry in range(0, 3):
             resp = http_req(*args, **kwargs)
             if resp.status == httpclient.SERVICE_UNAVAILABLE:
-                logger.warn("Storage service is not avaible temporaryly")
-                logger.info("Will retry later, in {0} seconds",
+                logger.warn("Storage service is temporarily unavailable. ")
+                logger.info("Will retry in {0} seconds. ",
                             LONG_WAITING_INTERVAL)
                 time.sleep(LONG_WAITING_INTERVAL)
             else:
                 return resp
-        raise ProtocolError(("Calling storage endpoint failed: {0}"
-                             "").format(resp.status))
+        raise ProtocolError(("Calling storage endpoint failed: "
+                             "{0}").format(resp.status))
 
     def fetch_manifest(self, version_uris):
         for version_uri in version_uris:
             logger.verbose("Fetch ext handler manifest: {0}", version_uri.uri)
             try:
-                resp = self.call_storage_service(restutil.http_get,
-                                                 version_uri.uri, None,
-                                                 chk_proxy=True)
+                resp = self.call_storage_service(
+                    restutil.http_get,
+                    version_uri.uri,
+                    None)
             except HttpError as e:
                 raise ProtocolError(ustr(e))
 
@@ -693,14 +718,16 @@ class WireClient(object):
 
     def get_hosting_env(self):
         if (self.hosting_env is None):
-            local_file = os.path.join(conf.get_lib_dir(), HOSTING_ENV_FILE_NAME)
+            local_file = os.path.join(conf.get_lib_dir(),
+                                      HOSTING_ENV_FILE_NAME)
             xml_text = self.fetch_cache(local_file)
             self.hosting_env = HostingEnv(xml_text)
         return self.hosting_env
 
     def get_shared_conf(self):
         if (self.shared_conf is None):
-            local_file = os.path.join(conf.get_lib_dir(), SHARED_CONF_FILE_NAME)
+            local_file = os.path.join(conf.get_lib_dir(),
+                                      SHARED_CONF_FILE_NAME)
             xml_text = self.fetch_cache(local_file)
             self.shared_conf = SharedConfig(xml_text)
         return self.shared_conf
@@ -796,7 +823,8 @@ class WireClient(object):
         headers = self.get_header_for_xml_content()
         try:
             resp = self.call_wireserver(restutil.http_post, health_report_uri,
-                                        health_report, headers=headers, max_retry=8)
+                                        health_report, headers=headers,
+                                        max_retry=8)
         except HttpError as e:
             raise ProtocolError((u"Failed to send provision status: {0}"
                                  u"").format(e))
@@ -820,7 +848,8 @@ class WireClient(object):
 
         if resp.status != httpclient.OK:
             logger.verbose(resp.read())
-            raise ProtocolError("Failed to send events:{0}".format(resp.status))
+            raise ProtocolError(
+                "Failed to send events:{0}".format(resp.status))
 
     def report_event(self, event_list):
         buf = {}
@@ -867,6 +896,7 @@ class WireClient(object):
             "x-ms-guest-agent-public-x509-cert": cert
         }
 
+
 class VersionInfo(object):
     def __init__(self, xml_text):
         """
@@ -880,14 +910,16 @@ class VersionInfo(object):
         xml_doc = parse_doc(xml_text)
         preferred = find(xml_doc, "Preferred")
         self.preferred = findtext(preferred, "Version")
-        logger.info("Fabric preferred wire protocol version:{0}", self.preferred)
+        logger.info("Fabric preferred wire protocol version:{0}",
+                    self.preferred)
 
         self.supported = []
         supported = find(xml_doc, "Supported")
         supported_version = findall(supported, "Version")
         for node in supported_version:
             version = gettext(node)
-            logger.verbose("Fabric supported wire protocol version:{0}", version)
+            logger.verbose("Fabric supported wire protocol version:{0}",
+                           version)
             self.supported.append(version)
 
     def get_preferred(self):
@@ -980,6 +1012,7 @@ class SharedConfig(object):
         # Not used currently
         return self
 
+
 class Certificates(object):
     """
     Object containing certificates of host and provisioned user.
@@ -1176,7 +1209,8 @@ class ExtensionsConfig(object):
             ext.sequenceNumber = seqNo
             ext.publicSettings = handler_settings.get("publicSettings")
             ext.protectedSettings = handler_settings.get("protectedSettings")
-            thumbprint = handler_settings.get("protectedSettingsCertThumbprint")
+            thumbprint = handler_settings.get(
+                "protectedSettingsCertThumbprint")
             ext.certificateThumbprint = thumbprint
             ext_handler.properties.extensions.append(ext)
 
@@ -1191,14 +1225,21 @@ class ExtensionManifest(object):
 
     def parse(self, xml_text):
         xml_doc = parse_doc(xml_text)
-        self._handle_packages(findall(find(xml_doc, "Plugins"), "Plugin"), False)
-        self._handle_packages(findall(find(xml_doc, "InternalPlugins"), "Plugin"), True)
+        self._handle_packages(findall(find(xml_doc,
+                                           "Plugins"),
+                                      "Plugin"),
+                              False)
+        self._handle_packages(findall(find(xml_doc,
+                                           "InternalPlugins"),
+                                      "Plugin"),
+                              True)
 
     def _handle_packages(self, packages, isinternal):
         for package in packages:
             version = findtext(package, "Version")
 
-            disallow_major_upgrade = findtext(package, "DisallowMajorVersionUpgrade")
+            disallow_major_upgrade = findtext(package,
+                                              "DisallowMajorVersionUpgrade")
             if disallow_major_upgrade is None:
                 disallow_major_upgrade = ''
             disallow_major_upgrade = disallow_major_upgrade.lower() == "true"
diff --git a/azurelinuxagent/common/utils/restutil.py b/azurelinuxagent/common/utils/restutil.py
index a7896507..ac022dfc 100644
--- a/azurelinuxagent/common/utils/restutil.py
+++ b/azurelinuxagent/common/utils/restutil.py
@@ -18,9 +18,7 @@
 #
 
 import time
-import platform
-import os
-import subprocess
+
 import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.exception import HttpError
@@ -32,6 +30,7 @@ REST api util functions
 
 RETRY_WAITING_INTERVAL = 10
 
+
 def _parse_url(url):
     o = urlparse(url)
     rel_uri = o.path
@@ -44,6 +43,7 @@ def _parse_url(url):
         secure = True
     return o.hostname, o.port, secure, rel_uri
 
+
 def get_http_proxy():
     """
     Get http_proxy and https_proxy from environment variables.
@@ -51,17 +51,19 @@ def get_http_proxy():
     """
     host = conf.get_httpproxy_host()
     port = conf.get_httpproxy_port()
-    return (host, port)
+    return host, port
+
 
 def _http_request(method, host, rel_uri, port=None, data=None, secure=False,
-                 headers=None, proxy_host=None, proxy_port=None):
+                  headers=None, proxy_host=None, proxy_port=None):
     url, conn = None, None
     if secure:
         port = 443 if port is None else port
         if proxy_host is not None and proxy_port is not None:
-            conn = httpclient.HTTPSConnection(proxy_host, proxy_port, timeout=10)
+            conn = httpclient.HTTPSConnection(proxy_host, proxy_port,
+                                              timeout=10)
             conn.set_tunnel(host, port)
-            #If proxy is used, full url is needed.
+            # If proxy is used, full url is needed.
             url = "https://{0}:{1}{2}".format(host, port, rel_uri)
         else:
             conn = httpclient.HTTPSConnection(host, port, timeout=10)
@@ -69,20 +71,23 @@ def _http_request(method, host, rel_uri, port=None, data=None, secure=False,
     else:
         port = 80 if port is None else port
         if proxy_host is not None and proxy_port is not None:
-            conn = httpclient.HTTPConnection(proxy_host, proxy_port, timeout=10)
-            #If proxy is used, full url is needed.
+            conn = httpclient.HTTPConnection(proxy_host, proxy_port,
+                                             timeout=10)
+            # If proxy is used, full url is needed.
             url = "http://{0}:{1}{2}".format(host, port, rel_uri)
         else:
             conn = httpclient.HTTPConnection(host, port, timeout=10)
             url = rel_uri
-    if headers == None:
+    if headers is None:
         conn.request(method, url, data)
     else:
         conn.request(method, url, data, headers)
     resp = conn.getresponse()
     return resp
 
-def http_request(method, url, data, headers=None, max_retry=3, chk_proxy=False):
+
+def http_request(method, url, data, headers=None, max_retry=3,
+                 chk_proxy=False):
     """
     Sending http request to server
     On error, sleep 10 and retry max_retry times.
@@ -92,28 +97,29 @@ def http_request(method, url, data, headers=None, max_retry=3, chk_proxy=False):
     logger.verbose("    Header={0}", headers)
     host, port, secure, rel_uri = _parse_url(url)
 
-    #Check proxy
+    # Check proxy
     proxy_host, proxy_port = (None, None)
     if chk_proxy:
         proxy_host, proxy_port = get_http_proxy()
 
-    #If httplib module is not built with ssl support. Fallback to http
+    # If httplib module is not built with ssl support. Fallback to http
     if secure and not hasattr(httpclient, "HTTPSConnection"):
         logger.warn("httplib is not built with ssl support")
         secure = False
 
-    #If httplib module doesn't support https tunnelling. Fallback to http
+    # If httplib module doesn't support https tunnelling. Fallback to http
     if secure and \
-            proxy_host is not None and \
-            proxy_port is not None and \
+                    proxy_host is not None and \
+                    proxy_port is not None and \
             not hasattr(httpclient.HTTPSConnection, "set_tunnel"):
-        logger.warn("httplib doesn't support https tunnelling(new in python 2.7)")
+        logger.warn("httplib does not support https tunnelling "
+                    "(new in python 2.7)")
         secure = False
 
     for retry in range(0, max_retry):
         try:
-            resp = _http_request(method, host, rel_uri, port=port, data=data, 
-                                 secure=secure, headers=headers, 
+            resp = _http_request(method, host, rel_uri, port=port, data=data,
+                                 secure=secure, headers=headers,
                                  proxy_host=proxy_host, proxy_port=proxy_port)
             logger.verbose("HTTP Resp: Status={0}", resp.status)
             logger.verbose("    Header={0}", resp.getheaders())
@@ -126,31 +132,36 @@ def http_request(method, url, data, headers=None, max_retry=3, chk_proxy=False):
         if retry < max_retry - 1:
             logger.info("Retry={0}, {1} {2}", retry, method, url)
             time.sleep(RETRY_WAITING_INTERVAL)
-    
+
     if url is not None and len(url) > 100:
-        url_log = url[0: 100] #In case the url is too long
+        url_log = url[0: 100]  # In case the url is too long
     else:
         url_log = url
     raise HttpError("HTTP Err: {0} {1}".format(method, url_log))
 
+
 def http_get(url, headers=None, max_retry=3, chk_proxy=False):
-    return http_request("GET", url, data=None, headers=headers, 
+    return http_request("GET", url, data=None, headers=headers,
                         max_retry=max_retry, chk_proxy=chk_proxy)
 
+
 def http_head(url, headers=None, max_retry=3, chk_proxy=False):
-    return http_request("HEAD", url, None, headers=headers, 
+    return http_request("HEAD", url, None, headers=headers,
                         max_retry=max_retry, chk_proxy=chk_proxy)
 
+
 def http_post(url, data, headers=None, max_retry=3, chk_proxy=False):
-    return http_request("POST", url, data, headers=headers, 
+    return http_request("POST", url, data, headers=headers,
                         max_retry=max_retry, chk_proxy=chk_proxy)
 
+
 def http_put(url, data, headers=None, max_retry=3, chk_proxy=False):
-    return http_request("PUT", url, data, headers=headers, 
+    return http_request("PUT", url, data, headers=headers,
                         max_retry=max_retry, chk_proxy=chk_proxy)
 
+
 def http_delete(url, headers=None, max_retry=3, chk_proxy=False):
-    return http_request("DELETE", url, None, headers=headers, 
+    return http_request("DELETE", url, None, headers=headers,
                         max_retry=max_retry, chk_proxy=chk_proxy)
 
-#End REST api util functions
+# End REST api util functions


Test patch:
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index bd3acafe..59af98b9 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -15,26 +15,16 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
-from tests.tools import *
+from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient
 from tests.protocol.mockwiredata import *
-import uuid
-import unittest
-import os
-import time
-from azurelinuxagent.common.utils.restutil import httpclient
-from azurelinuxagent.common.utils.cryptutil import CryptUtil
-from azurelinuxagent.common.protocol.restapi import *
-from azurelinuxagent.common.protocol.wire import WireClient, WireProtocol, \
-                                          TRANSPORT_PRV_FILE_NAME, \
-                                          TRANSPORT_CERT_FILE_NAME
 
 data_with_bom = b'\xef\xbb\xbfhehe'
 
+
 @patch("time.sleep")
 @patch("azurelinuxagent.common.protocol.wire.CryptUtil")
 @patch("azurelinuxagent.common.protocol.wire.restutil")
 class TestWireProtocolGetters(AgentTestCase):
-    
     def _test_getters(self, test_data, mock_restutil, MockCryptUtil, _):
         mock_restutil.http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
@@ -47,9 +37,9 @@ class TestWireProtocolGetters(AgentTestCase):
         for ext_handler in ext_handlers.extHandlers:
             protocol.get_ext_handler_pkgs(ext_handler)
 
-        crt1 = os.path.join(self.tmp_dir, 
-                           '33B0ABCE4673538650971C10F7D7397E71561F35.crt')
-        crt2 = os.path.join(self.tmp_dir, 
+        crt1 = os.path.join(self.tmp_dir,
+                            '33B0ABCE4673538650971C10F7D7397E71561F35.crt')
+        crt2 = os.path.join(self.tmp_dir,
                             '4037FBF5F1F3014F99B5D6C7799E9B20E6871CB3.crt')
         prv2 = os.path.join(self.tmp_dir,
                             '4037FBF5F1F3014F99B5D6C7799E9B20E6871CB3.prv')
@@ -57,7 +47,7 @@ class TestWireProtocolGetters(AgentTestCase):
         self.assertTrue(os.path.isfile(crt1))
         self.assertTrue(os.path.isfile(crt2))
         self.assertTrue(os.path.isfile(prv2))
-    
+
     def test_getters(self, *args):
         """Normal case"""
         test_data = WireProtocolData(DATA_FILE)
@@ -72,11 +62,42 @@ class TestWireProtocolGetters(AgentTestCase):
         """Extensions without any settings"""
         test_data = WireProtocolData(DATA_FILE_EXT_NO_SETTINGS)
         self._test_getters(test_data, *args)
-        
+
     def test_getters_ext_no_public(self, *args):
         """Extensions without any public settings"""
         test_data = WireProtocolData(DATA_FILE_EXT_NO_PUBLIC)
         self._test_getters(test_data, *args)
 
+    def test_call_storage_kwargs(self,
+                                 mock_restutil,
+                                 mock_cryptutil,
+                                 mock_sleep):
+        from azurelinuxagent.common.utils import restutil
+        with patch.object(restutil, 'http_get') as http_patch:
+            http_req = restutil.http_get
+            url = 'http://foo'
+            headers = {}
+
+            # no kwargs
+            WireClient.call_storage_service(http_req)
+            # kwargs, no chk_proxy
+            WireClient.call_storage_service(http_req,
+                                            url,
+                                            headers)
+            # kwargs, chk_proxy False
+            WireClient.call_storage_service(http_req,
+                                            url,
+                                            headers,
+                                            chk_proxy=False)
+            # kwargs, chk_proxy True
+            WireClient.call_storage_service(http_req,
+                                            url,
+                                            headers,
+                                            chk_proxy=True)
+            # assert
+            self.assertTrue(http_patch.call_count == 4)
+            for c in http_patch.call_args_list:
+                self.assertTrue(c[-1]['chk_proxy'] == True)
+
 if __name__ == '__main__':
     unittest.main()