diff --git a/custom_components/noonlight/switch.py b/custom_components/noonlight/switch.py
index 289495a..eed0f8d 100644
--- a/custom_components/noonlight/switch.py
+++ b/custom_components/noonlight/switch.py
@@ -7,9 +7,9 @@
 
 from homeassistant.components import persistent_notification
 try:
-    from homeassistant.components.switch import SwitchDevice
+    from homeassistant.components.switch import SwitchEntity
 except ImportError:
-    from homeassistant.components.switch import SwitchEntity as SwitchDevice
+    from homeassistant.components.switch import SwitchDevice as SwitchEntity
 from homeassistant.helpers.event import async_track_time_interval
 
 from . import (DOMAIN, EVENT_NOONLIGHT_TOKEN_REFRESHED,
@@ -37,7 +37,7 @@ def noonlight_token_refreshed():
         EVENT_NOONLIGHT_TOKEN_REFRESHED, noonlight_token_refreshed)
 
 
-class NoonlightSwitch(SwitchDevice):
+class NoonlightSwitch(SwitchEntity):
     """Representation of a Noonlight alarm switch."""
 
     def __init__(self, noonlight_integration):