From 6bb6a6ebe93ef6e053eb3ec638639b664a792f16 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus.schoutsen@mycase.com>
Date: Tue, 9 Aug 2016 20:23:57 -0700
Subject: [PATCH] Add notify demo data test

---
 homeassistant/components/notify/__init__.py | 11 +++++++----
 homeassistant/components/notify/demo.py     |  6 +++---
 tests/components/notify/test_demo.py        | 14 ++++++++++++++
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/notify/__init__.py b/homeassistant/components/notify/__init__.py
index 5b1c103a1bf..4b73c46b198 100644
--- a/homeassistant/components/notify/__init__.py
+++ b/homeassistant/components/notify/__init__.py
@@ -44,16 +44,19 @@ NOTIFY_SERVICE_SCHEMA = vol.Schema({
 _LOGGER = logging.getLogger(__name__)
 
 
-def send_message(hass, message, title=None):
+def send_message(hass, message, title=None, data=None):
     """Send a notification message."""
-    data = {
+    info = {
         ATTR_MESSAGE: message
     }
 
     if title is not None:
-        data[ATTR_TITLE] = title
+        info[ATTR_TITLE] = title
 
-    hass.services.call(DOMAIN, SERVICE_NOTIFY, data)
+    if data is not None:
+        info[ATTR_DATA] = data
+
+    hass.services.call(DOMAIN, SERVICE_NOTIFY, info)
 
 
 def setup(hass, config):
diff --git a/homeassistant/components/notify/demo.py b/homeassistant/components/notify/demo.py
index c051bce020d..4685b90e880 100644
--- a/homeassistant/components/notify/demo.py
+++ b/homeassistant/components/notify/demo.py
@@ -4,7 +4,7 @@ Demo notification service.
 For more details about this platform, please refer to the documentation
 https://home-assistant.io/components/demo/
 """
-from homeassistant.components.notify import ATTR_TITLE, BaseNotificationService
+from homeassistant.components.notify import BaseNotificationService
 
 EVENT_NOTIFY = "notify"
 
@@ -24,5 +24,5 @@ class DemoNotificationService(BaseNotificationService):
 
     def send_message(self, message="", **kwargs):
         """Send a message to a user."""
-        title = kwargs.get(ATTR_TITLE)
-        self.hass.bus.fire(EVENT_NOTIFY, {"title": title, "message": message})
+        kwargs['message'] = message
+        self.hass.bus.fire(EVENT_NOTIFY, kwargs)
diff --git a/tests/components/notify/test_demo.py b/tests/components/notify/test_demo.py
index 0d4f2115ca7..079daa48498 100644
--- a/tests/components/notify/test_demo.py
+++ b/tests/components/notify/test_demo.py
@@ -45,3 +45,17 @@ class TestNotifyDemo(unittest.TestCase):
         last_event = self.events[-1]
         self.assertEqual(last_event.data[notify.ATTR_TITLE], 'temperature')
         self.assertEqual(last_event.data[notify.ATTR_MESSAGE], '10')
+
+    def test_method_forwards_correct_data(self):
+        """Test that all data from the service gets forwarded to service."""
+        notify.send_message(self.hass, 'my message', 'my title',
+                            {'hello': 'world'})
+        self.hass.pool.block_till_done()
+        self.assertTrue(len(self.events) == 1)
+        data = self.events[0].data
+        assert {
+            'message': 'my message',
+            'target': None,
+            'title': 'my title',
+            'data': {'hello': 'world'}
+        } == data