diff --git a/__init__.py b/__init__.py
index 72e0e3e..186608a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -12,7 +12,8 @@
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import ToggleEntity
 from homeassistant.helpers.entity_component import EntityComponent
-from homeassistant.helpers.config_validation import PLATFORM_SCHEMA  # noqa
+from homeassistant.helpers.config_validation import (  # noqa
+    PLATFORM_SCHEMA, PLATFORM_SCHEMA_BASE)
 from homeassistant.loader import bind_hass
 
 _LOGGER = logging.getLogger(__name__)