Repository: Azure/WALinuxAgent
Pull request #444 (2016-09-23 20:04:51)

Send additional headers when using HostGAPlugin (#443)
Taking care of "Send additional headers when using HostGAPlugin  #443"
Added necessary headers when sending put vm status request to HostGAPlugin.

Added UTs


Patch:
diff --git a/azurelinuxagent/common/protocol/hostplugin.py b/azurelinuxagent/common/protocol/hostplugin.py
index 66d58898..1bfee265 100644
--- a/azurelinuxagent/common/protocol/hostplugin.py
+++ b/azurelinuxagent/common/protocol/hostplugin.py
@@ -28,13 +28,14 @@ API_VERSION = "2015-09-01"
 
 
 class HostPluginProtocol(object):
-    def __init__(self, endpoint):
+    def __init__(self, endpoint, goal_state):
         if endpoint is None:
             raise ProtocolError("Host plugin endpoint not provided")
         self.is_initialized = False
         self.is_available = False
         self.api_versions = None
         self.endpoint = endpoint
+        self.goal_state = goal_state
 
     def ensure_initialized(self):
         if not self.is_initialized:
@@ -48,7 +49,8 @@ class HostPluginProtocol(object):
                                                  HOST_PLUGIN_PORT)
         logger.info("getting API versions at [{0}]".format(url))
         try:
-            response = restutil.http_get(url)
+            headers = { "x-ms-containerid": self.goal_state.container_id }
+            response = restutil.http_get(url, headers)
             if response.status != httpclient.OK:
                 logger.error(
                     "get API versions returned status code [{0}]".format(
@@ -73,8 +75,11 @@ class HostPluginProtocol(object):
             return
         url = URI_FORMAT_PUT_VM_STATUS.format(self.endpoint, HOST_PLUGIN_PORT)
         logger.verbose("Posting VM status to host channel [{0}]".format(url))
-        status = textutil.b64encode(status_blob.vm_status.vmAgent.status)
-        headers = {"x-ms-version": API_VERSION}
+        status = textutil.b64encode(status_blob.data)
+        headers = {"x-ms-version": API_VERSION,
+                   "Content-type": "application/json",
+                   "x-ms-containerid": self.goal_state.container_id,
+                   "x-ms-host-config-name": self.goal_state.role_instance_config_name}
         blob_headers = [{'headerName': 'x-ms-version',
                          'headerValue': status_blob.__storage_version__},
                         {'headerName': 'x-ms-blob-type',
@@ -89,7 +94,7 @@ class HostPluginProtocol(object):
         except HttpError as e:
             logger.error("put VM status failed with [{0}]".format(e))
 
-    def put_vm_log(self, content, container_id, deployment_id):
+    def put_vm_log(self, content):
         """
         Try to upload the given content to the host plugin
         :param deployment_id: the deployment id, which is obtained from the
@@ -102,18 +107,18 @@ class HostPluginProtocol(object):
         if not self.ensure_initialized():
             logger.error("host plugin channel is not available")
             return
-        if content is None or container_id is None or deployment_id is None:
+        if content is None or self.goal_state.container_id is None or self.goal_state.deployment_id is None:
             logger.error(
                 "invalid arguments passed: "
                 "[{0}], [{1}], [{2}]".format(
                     content,
-                    container_id,
-                    deployment_id))
+                    self.goal_state.container_id,
+                    self.goal_state.deployment_id))
             return
         url = URI_FORMAT_PUT_LOG.format(self.endpoint, HOST_PLUGIN_PORT)
 
-        headers = {"x-ms-vmagentlog-deploymentid": deployment_id,
-                   "x-ms-vmagentlog-containerid": container_id}
+        headers = {"x-ms-vmagentlog-deploymentid": self.goal_state.deployment_id,
+                   "x-ms-vmagentlog-containerid": self.goal_state.container_id}
         logger.info("put VM log at [{0}]".format(url))
         try:
             response = restutil.http_put(url, content, headers)
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index cb32bd4e..771f2c38 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -509,8 +509,8 @@ class WireClient(object):
         self.ext_conf = None
         self.last_request = 0
         self.req_count = 0
+        self.host_plugin = None
         self.status_blob = StatusBlob(self)
-        self.host_plugin = HostPluginProtocol(self.endpoint)
 
     def prevent_throttling(self):
         """
@@ -702,6 +702,8 @@ class WireClient(object):
                 self.update_shared_conf(goal_state)
                 self.update_certs(goal_state)
                 self.update_ext_conf(goal_state)
+                if self.host_plugin is not None:
+                    self.host_plugin.goal_state = goal_state
                 return
             except WireProtocolResourceGone:
                 logger.info("Incarnation is out of date. Update goalstate.")
@@ -795,8 +797,7 @@ class WireClient(object):
         ext_conf = self.get_ext_conf()
         if ext_conf.status_upload_blob is not None:
             if not self.status_blob.upload(ext_conf.status_upload_blob):
-                self.host_plugin.put_vm_status(self.status_blob,
-                                               ext_conf.status_upload_blob)
+                self.get_host_plugin().put_vm_status(self.status_blob, ext_conf.status_upload_blob)
 
     def report_role_prop(self, thumbprint):
         goal_state = self.get_goal_state()
@@ -911,6 +912,10 @@ class WireClient(object):
             "x-ms-guest-agent-public-x509-cert": cert
         }
 
+    def get_host_plugin(self):
+        if self.host_plugin is None:
+            self.host_plugin = HostPluginProtocol(self.endpoint, self.get_goal_state())
+        return self.host_plugin
 
 class VersionInfo(object):
     def __init__(self, xml_text):
@@ -956,6 +961,7 @@ class GoalState(object):
         self.certs_uri = None
         self.ext_uri = None
         self.role_instance_id = None
+        self.role_instance_config_name = None
         self.container_id = None
         self.load_balancer_probe_port = None
         self.parse(xml_text)
@@ -974,6 +980,8 @@ class GoalState(object):
         self.ext_uri = findtext(xml_doc, "ExtensionsConfig")
         role_instance = find(xml_doc, "RoleInstance")
         self.role_instance_id = findtext(role_instance, "InstanceId")
+        role_config = find(role_instance, "Configuration")
+        self.role_instance_config_name = findtext(role_config, "ConfigName")
         container = find(xml_doc, "Container")
         self.container_id = findtext(container, "ContainerId")
         lbprobe_ports = find(xml_doc, "LBProbePorts")
diff --git a/azurelinuxagent/common/version.py b/azurelinuxagent/common/version.py
index ef278901..4018670f 100644
--- a/azurelinuxagent/common/version.py
+++ b/azurelinuxagent/common/version.py
@@ -51,7 +51,7 @@ def get_distro():
 
 AGENT_NAME = "WALinuxAgent"
 AGENT_LONG_NAME = "Azure Linux Agent"
-AGENT_VERSION = '2.1.6.7'
+AGENT_VERSION = '2.1.6.6'
 AGENT_LONG_VERSION = "{0}-{1}".format(AGENT_NAME, AGENT_VERSION)
 AGENT_DESCRIPTION = """\
 The Azure Linux Agent supports the provisioning and running of Linux


Test patch:
diff --git a/tests/data/wire/goal_state.xml b/tests/data/wire/goal_state.xml
index 960444bb..0fd144f2 100644
--- a/tests/data/wire/goal_state.xml
+++ b/tests/data/wire/goal_state.xml
@@ -20,6 +20,7 @@
          <Certificates>http://certificatesuri/</Certificates>
          <ExtensionsConfig>http://extensionsconfiguri/</ExtensionsConfig>
          <FullConfig>http://fullconfiguri/</FullConfig>
+         <ConfigName>DummyRoleConfigName.xml</ConfigName>
          </Configuration>
        </RoleInstance>
      </RoleInstanceList>
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index 296a1f58..b0a03c52 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -19,6 +19,7 @@ import unittest
 
 import azurelinuxagent.common.protocol.restapi as restapi
 import azurelinuxagent.common.protocol.wire as wire
+from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE
 from tests.tools import *
 
 wireserver_url = "168.63.129.16"
@@ -28,11 +29,16 @@ api_versions = '["2015-09-01"]'
 
 class TestHostPlugin(AgentTestCase):
     def test_fallback(self):
-        with patch.object(wire.HostPluginProtocol,
-                          "put_vm_status") as patch_put:
-            with patch.object(wire.StatusBlob, "upload") as patch_upload:
-                patch_upload.return_value = False
+        """
+        Validate fallback to upload status using HostGAPlugin is happening when status reporting via
+        default method is unsuccessful
+        """
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
+
+        with patch.object(wire.HostPluginProtocol, "put_vm_status") as patch_put:
+            with patch.object(wire.StatusBlob, "upload", return_value=False) as patch_upload:
                 wire_protocol_client = wire.WireProtocol(wireserver_url).client
+                wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
                 wire_protocol_client.ext_conf = wire.ExtensionsConfig(None)
                 wire_protocol_client.ext_conf.status_upload_blob = sas_url
                 wire_protocol_client.upload_status_blob()
@@ -40,22 +46,26 @@ class TestHostPlugin(AgentTestCase):
                                 "Fallback was not engaged")
                 self.assertTrue(patch_put.call_args[0][1] == sas_url)
 
-    def test_fallback_data(self):
+    def test_validate_http_request(self):
+        """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
         from azurelinuxagent.common.protocol.hostplugin import API_VERSION
         from azurelinuxagent.common.utils import restutil
         exp_method = 'PUT'
         exp_url = 'http://{0}:32526/status'.format(wireserver_url)
-        exp_data = '{"content": "UmVhZHk=", "headers": [{"headerName": ' \
+        exp_data = '{"content": "eyJkdW1teSI6ICJkYXRhIn0=", "headers": [{"headerName": ' \
                    '"x-ms-version", "headerValue": "2014-02-14"}, ' \
                    '{"headerName": "x-ms-blob-type", "headerValue": null}], ' \
                    '"requestUri": "http://sas_url"}'
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
+
         with patch.object(restutil, "http_request") as patch_http:
             wire_protocol_client = wire.WireProtocol(wireserver_url).client
-            plugin = wire_protocol_client.host_plugin
+            wire_protocol_client.get_goal_state = Mock(return_value=test_goal_state)
+            plugin = wire_protocol_client.get_host_plugin()
             blob = wire_protocol_client.status_blob
             blob.vm_status = restapi.VMStatus()
-            blob.vm_status.vmAgent.status = 'Ready'
-            with patch.object(plugin, "get_api_versions") as patch_api:
+            blob.data = '{"dummy": "data"}'
+            with patch.object(plugin, 'get_api_versions') as patch_api:
                 patch_api.return_value = API_VERSION
                 plugin.put_vm_status(blob, sas_url)
                 self.assertTrue(patch_http.call_count == 1)
@@ -63,7 +73,16 @@ class TestHostPlugin(AgentTestCase):
                 self.assertTrue(patch_http.call_args[0][1] == exp_url)
                 self.assertTrue(patch_http.call_args[0][2] == exp_data)
 
+                # Assert headers
+                headers = patch_http.call_args[1]['headers']
+                self.assertEqual(headers['x-ms-containerid'], test_goal_state.container_id)
+                self.assertEqual(headers['x-ms-host-config-name'], test_goal_state.role_instance_config_name)
+
     def test_no_fallback(self):
+        """
+        Validate fallback to upload status using HostGAPlugin is not happening when status reporting via
+        default method is successful
+        """
         with patch.object(wire.HostPluginProtocol,
                           "put_vm_status") as patch_put:
             with patch.object(wire.StatusBlob, "upload") as patch_upload:
@@ -75,22 +94,27 @@ class TestHostPlugin(AgentTestCase):
                 self.assertTrue(patch_put.call_count == 0,
                                 "Fallback was engaged")
 
-    def test_init_put(self):
+    def test_validate_http_put(self):
+        """Validate correct set of data is sent to HostGAPlugin when reporting VM status"""
+        test_goal_state = wire.GoalState(WireProtocolData(DATA_FILE).goal_state)
         expected_url = "http://168.63.129.16:32526/status"
-        expected_headers = {'x-ms-version': '2015-09-01'}
-        expected_content = '{"content": "b2s=", ' \
+        expected_headers = {'x-ms-version': '2015-09-01',
+                            "Content-type": "application/json",
+                            "x-ms-containerid": test_goal_state.container_id,
+                            "x-ms-host-config-name": test_goal_state.role_instance_config_name}
+        expected_content = '{"content": "eyJkdW1teSI6ICJkYXRhIn0=", ' \
                            '"headers": [{"headerName": "x-ms-version", ' \
                            '"headerValue": "2014-02-14"}, ' \
                            '{"headerName": "x-ms-blob-type", "headerValue": ' \
                            '"BlockBlob"}], ' \
                            '"requestUri": "http://sas_url"}'
 
-        host_client = wire.HostPluginProtocol(wireserver_url)
+        host_client = wire.HostPluginProtocol(wireserver_url, test_goal_state)
         self.assertFalse(host_client.is_initialized)
         self.assertTrue(host_client.api_versions is None)
         status_blob = wire.StatusBlob(None)
         status_blob.vm_status = restapi.VMStatus()
-        status_blob.vm_status.vmAgent.status = 'ok'
+        status_blob.data = '{"dummy": "data"}'
         status_blob.type = "BlockBlob"
         with patch.object(wire.HostPluginProtocol,
                           "get_api_versions") as patch_get:
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index 59af98b9..5a29ebbf 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -15,11 +15,12 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
-from azurelinuxagent.common.protocol.wire import WireProtocol, WireClient
+from azurelinuxagent.common.protocol.wire import *
 from tests.protocol.mockwiredata import *
 
 data_with_bom = b'\xef\xbb\xbfhehe'
-
+testurl = 'http://foo'
+wireserver_url = '168.63.129.16'
 
 @patch("time.sleep")
 @patch("azurelinuxagent.common.protocol.wire.CryptUtil")
@@ -29,7 +30,7 @@ class TestWireProtocolGetters(AgentTestCase):
         mock_restutil.http_get.side_effect = test_data.mock_http_get
         MockCryptUtil.side_effect = test_data.mock_crypt_util
 
-        protocol = WireProtocol("foo.bar")
+        protocol = WireProtocol(wireserver_url)
         protocol.detect()
         protocol.get_vminfo()
         protocol.get_certs()
@@ -75,7 +76,7 @@ class TestWireProtocolGetters(AgentTestCase):
         from azurelinuxagent.common.utils import restutil
         with patch.object(restutil, 'http_get') as http_patch:
             http_req = restutil.http_get
-            url = 'http://foo'
+            url = testurl
             headers = {}
 
             # no kwargs
@@ -99,5 +100,44 @@ class TestWireProtocolGetters(AgentTestCase):
             for c in http_patch.call_args_list:
                 self.assertTrue(c[-1]['chk_proxy'] == True)
 
+    def test_get_host_ga_plugin(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+
+        with patch.object(WireClient, "get_goal_state", return_value = goal_state) as patch_get_goal_state:
+            host_plugin = wire_protocol_client.get_host_plugin()
+            self.assertEqual(goal_state, host_plugin.goal_state)
+            patch_get_goal_state.assert_called_once()
+
+    def test_upload_status_blob_default(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.status_upload_blob = testurl
+
+        with patch.object(WireClient, "get_goal_state") as patch_get_goal_state:
+            with patch.object(HostPluginProtocol, "put_vm_status") as patch_host_ga_plugin_upload:
+                with patch.object(StatusBlob, "upload", return_value = True) as patch_default_upload:
+                    wire_protocol_client.upload_status_blob()
+
+                    patch_default_upload.assert_called_once_with(testurl)
+                    patch_get_goal_state.assert_not_called()
+                    patch_host_ga_plugin_upload.assert_not_called()
+
+    def test_upload_status_blob_host_ga_plugin(self, *args):
+        wire_protocol_client = WireProtocol(wireserver_url).client
+        wire_protocol_client.ext_conf = ExtensionsConfig(None)
+        wire_protocol_client.ext_conf.status_upload_blob = testurl
+        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+
+        with patch.object(HostPluginProtocol, "put_vm_status") as patch_host_ga_plugin_upload:
+            with patch.object(StatusBlob, "upload", return_value=False) as patch_default_upload:
+                wire_protocol_client.get_goal_state = Mock(return_value = goal_state)
+                wire_protocol_client.upload_status_blob()
+
+                patch_default_upload.assert_called_once_with(testurl)
+                wire_protocol_client.get_goal_state.assert_called_once()
+                patch_host_ga_plugin_upload.assert_called_once_with(wire_protocol_client.status_blob, testurl)
+
+
 if __name__ == '__main__':
     unittest.main()