From 226066eafd02063acbc04ea758568d3cd255b0f4 Mon Sep 17 00:00:00 2001
From: Mike Megally <cmsimike@users.noreply.github.com>
Date: Tue, 18 Apr 2017 02:50:43 -0700
Subject: [PATCH] exposed content_type in rest_command (#7101)

* exposed content_type in rest_command, which allows for manually specifying the content_type for more-strict api endpoints

* fixed up column length

Length was 86 chars, and it needed to be 79

* double import of HTTP_HEADER_CONTENT_TYPE

Removed the accidental double-import of HTTP_HEADER_CONTENT_TYPE

* moved rest_command-specific config value into component

* if no content_type, default to None

* unit test

* newline

* unused CONTENT_TYPE_TEXT_PLAIN

* removed the http-agnostic abstraction hass provided in favor of aiohttps hdrs constant
---
 homeassistant/components/rest_command.py | 12 +++++++++++-
 tests/components/test_rest_command.py    | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/rest_command.py b/homeassistant/components/rest_command.py
index 7084339ab42..5bf629ed37f 100644
--- a/homeassistant/components/rest_command.py
+++ b/homeassistant/components/rest_command.py
@@ -8,6 +8,7 @@ import asyncio
 import logging
 
 import aiohttp
+from aiohttp import hdrs
 import async_timeout
 import voluptuous as vol
 
@@ -31,6 +32,8 @@ SUPPORT_REST_METHODS = [
     'delete',
 ]
 
+CONF_CONTENT_TYPE = 'content_type'
+
 COMMAND_SCHEMA = vol.Schema({
     vol.Required(CONF_URL): cv.template,
     vol.Optional(CONF_METHOD, default=DEFAULT_METHOD):
@@ -39,6 +42,7 @@ COMMAND_SCHEMA = vol.Schema({
     vol.Inclusive(CONF_PASSWORD, 'authentication'): cv.string,
     vol.Optional(CONF_PAYLOAD): cv.template,
     vol.Optional(CONF_TIMEOUT, default=DEFAULT_TIMEOUT): vol.Coerce(int),
+    vol.Optional(CONF_CONTENT_TYPE): cv.string
 })
 
 CONFIG_SCHEMA = vol.Schema({
@@ -72,6 +76,11 @@ def async_setup(hass, config):
             template_payload = command_config[CONF_PAYLOAD]
             template_payload.hass = hass
 
+        headers = None
+        if CONF_CONTENT_TYPE in command_config:
+            content_type = command_config[CONF_CONTENT_TYPE]
+            headers = {hdrs.CONTENT_TYPE: content_type}
+
         @asyncio.coroutine
         def async_service_handler(service):
             """Execute a shell command service."""
@@ -86,7 +95,8 @@ def async_setup(hass, config):
                     request = yield from getattr(websession, method)(
                         template_url.async_render(variables=service.data),
                         data=payload,
-                        auth=auth
+                        auth=auth,
+                        headers=headers
                     )
 
                 if request.status < 400:
diff --git a/tests/components/test_rest_command.py b/tests/components/test_rest_command.py
index 0648a30c922..9dbea53cd64 100644
--- a/tests/components/test_rest_command.py
+++ b/tests/components/test_rest_command.py
@@ -221,3 +221,22 @@ 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):
+        """Call a rest command with a content type."""
+        data = {
+            'payload': 'item',
+            'content_type': 'text/plain'
+        }
+        self.config[rc.DOMAIN]['post_test'].update(data)
+
+        with assert_setup_component(4):
+            setup_component(self.hass, rc.DOMAIN, self.config)
+
+        aioclient_mock.post(self.url, content=b'success')
+
+        self.hass.services.call(rc.DOMAIN, 'post_test', {})
+        self.hass.block_till_done()
+
+        assert len(aioclient_mock.mock_calls) == 1
+        assert aioclient_mock.mock_calls[0][2] == b'item'