Add custom header support for rest_command (#12646)

* Add support for specifying custom headers for rest_command.
* Added headers configuration to behave similarly to the rest sensor.
* Replaced test_rest_command_content_type which only validated the
  configuration with test_rest_command_headers which tests several
  combinations of parameters that affect the request headers.

homeassistant/components/rest_command.py

@@ -14,7 +14,7 @@ import voluptuous as vol
 
 from homeassistant.const import (
     CONF_TIMEOUT, CONF_USERNAME, CONF_PASSWORD, CONF_URL, CONF_PAYLOAD,
-    CONF_METHOD)
+    CONF_METHOD, CONF_HEADERS)
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 
@@ -38,6 +38,7 @@ COMMAND_SCHEMA = vol.Schema({
     vol.Required(CONF_URL): cv.template,
     vol.Optional(CONF_METHOD, default=DEFAULT_METHOD):
         vol.All(vol.Lower, vol.In(SUPPORT_REST_METHODS)),
+    vol.Optional(CONF_HEADERS): vol.Schema({cv.string: cv.string}),
     vol.Inclusive(CONF_USERNAME, 'authentication'): cv.string,
     vol.Inclusive(CONF_PASSWORD, 'authentication'): cv.string,
     vol.Optional(CONF_PAYLOAD): cv.template,
@@ -77,9 +78,14 @@ def async_setup(hass, config):
             template_payload.hass = hass
 
         headers = None
+        if CONF_HEADERS in command_config:
+            headers = command_config[CONF_HEADERS]
+
         if CONF_CONTENT_TYPE in command_config:
             content_type = command_config[CONF_CONTENT_TYPE]
-            headers = {hdrs.CONTENT_TYPE: content_type}
+            if headers is None:
+                headers = {}
+            headers[hdrs.CONTENT_TYPE] = content_type
 
         @asyncio.coroutine
         def async_service_handler(service):

homeassistant/components/sensor/rest.py

@@ -35,7 +35,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_RESOURCE): cv.url,
     vol.Optional(CONF_AUTHENTICATION):
         vol.In([HTTP_BASIC_AUTHENTICATION, HTTP_DIGEST_AUTHENTICATION]),
-    vol.Optional(CONF_HEADERS): {cv.string: cv.string},
+    vol.Optional(CONF_HEADERS): vol.Schema({cv.string: cv.string}),
     vol.Optional(CONF_JSON_ATTRS, default=[]): cv.ensure_list_csv,
     vol.Optional(CONF_METHOD, default=DEFAULT_METHOD): vol.In(METHODS),
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,

tests/components/test_rest_command.py

@@ -222,21 +222,82 @@ class TestRestCommandComponent(object):
         assert len(aioclient_mock.mock_calls) == 1
         assert aioclient_mock.mock_calls[0][2] == b'data'
 
-    def test_rest_command_content_type(self, aioclient_mock):
+    def test_rest_command_headers(self, aioclient_mock):
-        """Call a rest command with a content type."""
+        """Call a rest command with custom headers and content types."""
-        data = {
+        header_config_variations = {
-            'payload': 'item',
+            rc.DOMAIN: {
+                'no_headers_test': {},
+                'content_type_test': {
+                    'content_type': 'text/plain'
+                },
+                'headers_test': {
+                    'headers': {
+                        'Accept': 'application/json',
+                        'User-Agent': 'Mozilla/5.0'
+                    }
+                },
+                'headers_and_content_type_test': {
+                    'headers': {
+                        'Accept': 'application/json'
+                    },
+                    'content_type': 'text/plain'
+                },
+                'headers_and_content_type_override_test': {
+                    'headers': {
+                        'Accept': 'application/json',
+                        aiohttp.hdrs.CONTENT_TYPE: 'application/pdf'
+                    },
                     'content_type': 'text/plain'
                 }
-        self.config[rc.DOMAIN]['post_test'].update(data)
+            }
+        }
 
-        with assert_setup_component(4):
+        # add common parameters
-            setup_component(self.hass, rc.DOMAIN, self.config)
+        for variation in header_config_variations[rc.DOMAIN].values():
+            variation.update({'url': self.url, 'method': 'post',
+                              'payload': 'test data'})
 
+        with assert_setup_component(5):
+            setup_component(self.hass, rc.DOMAIN, header_config_variations)
+
+        # provide post request data
         aioclient_mock.post(self.url, content=b'success')
 
-        self.hass.services.call(rc.DOMAIN, 'post_test', {})
+        for test_service in ['no_headers_test',
-        self.hass.block_till_done()
+                             'content_type_test',
+                             'headers_test',
+                             'headers_and_content_type_test',
+                             'headers_and_content_type_override_test']:
+            self.hass.services.call(rc.DOMAIN, test_service, {})
 
-        assert len(aioclient_mock.mock_calls) == 1
+        self.hass.block_till_done()
-        assert aioclient_mock.mock_calls[0][2] == b'item'
+        assert len(aioclient_mock.mock_calls) == 5
+
+        # no_headers_test
+        assert aioclient_mock.mock_calls[0][3] is None
+
+        # content_type_test
+        assert len(aioclient_mock.mock_calls[1][3]) == 1
+        assert aioclient_mock.mock_calls[1][3].get(
+            aiohttp.hdrs.CONTENT_TYPE) == 'text/plain'
+
+        # headers_test
+        assert len(aioclient_mock.mock_calls[2][3]) == 2
+        assert aioclient_mock.mock_calls[2][3].get(
+            'Accept') == 'application/json'
+        assert aioclient_mock.mock_calls[2][3].get(
+            'User-Agent') == 'Mozilla/5.0'
+
+        # headers_and_content_type_test
+        assert len(aioclient_mock.mock_calls[3][3]) == 2
+        assert aioclient_mock.mock_calls[3][3].get(
+            aiohttp.hdrs.CONTENT_TYPE) == 'text/plain'
+        assert aioclient_mock.mock_calls[3][3].get(
+            'Accept') == 'application/json'
+
+        # headers_and_content_type_override_test
+        assert len(aioclient_mock.mock_calls[4][3]) == 2
+        assert aioclient_mock.mock_calls[4][3].get(
+            aiohttp.hdrs.CONTENT_TYPE) == 'text/plain'
+        assert aioclient_mock.mock_calls[4][3].get(
+            'Accept') == 'application/json'