diff --git a/homeassistant/components/demo.py b/homeassistant/components/demo.py
index 67127f7e3cd..961feea705e 100644
--- a/homeassistant/components/demo.py
+++ b/homeassistant/components/demo.py
@@ -8,11 +8,11 @@ import random
 
 import homeassistant as ha
 import homeassistant.loader as loader
-from homeassistant.components import (SERVICE_TURN_ON, SERVICE_TURN_OFF,
-                                      STATE_ON, STATE_OFF, ATTR_ENTITY_PICTURE,
-                                      extract_entity_ids)
-from homeassistant.components.light import (ATTR_XY_COLOR, ATTR_BRIGHTNESS,
-                                            GROUP_NAME_ALL_LIGHTS)
+from homeassistant.components import (
+    SERVICE_TURN_ON, SERVICE_TURN_OFF, STATE_ON, STATE_OFF,
+    ATTR_ENTITY_PICTURE, ATTR_ENTITY_ID, extract_entity_ids)
+from homeassistant.components.light import (
+    ATTR_XY_COLOR, ATTR_BRIGHTNESS, GROUP_NAME_ALL_LIGHTS)
 from homeassistant.util import split_entity_id
 
 DOMAIN = "demo"
@@ -24,6 +24,9 @@ def setup(hass, config):
     """ Setup a demo environment. """
     group = loader.get_component('group')
 
+    config.setdefault(ha.DOMAIN, {})
+    config.setdefault(DOMAIN, {})
+
     if config[DOMAIN].get('hide_demo_state') != '1':
         hass.states.set('a.Demo_Mode', 'Enabled')
 
@@ -35,7 +38,12 @@ def setup(hass, config):
 
     def mock_turn_on(service):
         """ Will fake the component has been turned on. """
-        for entity_id in extract_entity_ids(hass, service):
+        if service.data and ATTR_ENTITY_ID in service.data:
+            entity_ids = extract_entity_ids(hass, service)
+        else:
+            entity_ids = hass.get_entity_ids(service.domain)
+
+        for entity_id in entity_ids:
             domain, _ = split_entity_id(entity_id)
 
             if domain == "light":
@@ -48,7 +56,12 @@ def setup(hass, config):
 
     def mock_turn_off(service):
         """ Will fake the component has been turned off. """
-        for entity_id in extract_entity_ids(hass, service):
+        if service.data and ATTR_ENTITY_ID in service.data:
+            entity_ids = extract_entity_ids(hass, service)
+        else:
+            entity_ids = hass.get_entity_ids(service.domain)
+
+        for entity_id in entity_ids:
             hass.states.set(entity_id, STATE_OFF)
 
     # Setup sun
diff --git a/test/test_component_demo.py b/test/test_component_demo.py
new file mode 100644
index 00000000000..901d62b4f60
--- /dev/null
+++ b/test/test_component_demo.py
@@ -0,0 +1,77 @@
+"""
+test.test_component_demo
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Tests demo component.
+"""
+# pylint: disable=too-many-public-methods,protected-access
+import unittest
+import datetime as dt
+
+import ephem
+
+import homeassistant as ha
+import homeassistant.components.demo as demo
+from homeassistant.components import (
+    SERVICE_TURN_ON, SERVICE_TURN_OFF, STATE_ON, STATE_OFF, ATTR_ENTITY_ID)
+
+
+class TestDemo(unittest.TestCase):
+    """ Test the demo module. """
+
+    def setUp(self):  # pylint: disable=invalid-name
+        self.hass = ha.HomeAssistant()
+
+    def tearDown(self):  # pylint: disable=invalid-name
+        """ Stop down stuff we started. """
+        self.hass.stop()
+
+    def test_services(self):
+        """ Test the demo services. """
+        # Test turning on and off different types
+        demo.setup(self.hass, {})
+
+        for domain in ('light', 'switch'):
+            # Focus on 1 entity
+            entity_id = self.hass.get_entity_ids(domain)[0]
+
+            self.hass.call_service(
+                domain, SERVICE_TURN_ON, {ATTR_ENTITY_ID: entity_id})
+
+            self.hass._pool.block_till_done()
+
+            self.assertEqual(STATE_ON, self.hass.states.get(entity_id).state)
+
+            self.hass.call_service(
+                domain, SERVICE_TURN_OFF, {ATTR_ENTITY_ID: entity_id})
+
+            self.hass._pool.block_till_done()
+
+            self.assertEqual(STATE_OFF, self.hass.states.get(entity_id).state)
+
+            # Act on all
+            self.hass.call_service(domain, SERVICE_TURN_ON)
+
+            self.hass._pool.block_till_done()
+
+            for entity_id in self.hass.get_entity_ids(domain):
+                self.assertEqual(
+                    STATE_ON, self.hass.states.get(entity_id).state)
+
+            self.hass.call_service(domain, SERVICE_TURN_OFF)
+
+            self.hass._pool.block_till_done()
+
+            for entity_id in self.hass.get_entity_ids(domain):
+                self.assertEqual(
+                    STATE_OFF, self.hass.states.get(entity_id).state)
+
+    def test_hiding_demo_state(self):
+        """ Test if you can hide the demo card. """
+        demo.setup(self.hass, {demo.DOMAIN: {'hide_demo_state': '1'}})
+
+        self.assertIsNone(self.hass.states.get('a.Demo_Mode'))
+
+        demo.setup(self.hass, {demo.DOMAIN: {'hide_demo_state': '0'}})
+
+        self.assertIsNotNone(self.hass.states.get('a.Demo_Mode'))