Merge pull request #33118 from home-assistant/rc

0.107.5

homeassistant/components/huawei_lte/manifest.json

@@ -5,7 +5,7 @@
   "documentation": "https://www.home-assistant.io/integrations/huawei_lte",
   "requirements": [
     "getmac==0.8.1",
-    "huawei-lte-api==1.4.10",
+    "huawei-lte-api==1.4.11",
     "stringcase==1.2.0",
     "url-normalize==1.4.1"
   ],

homeassistant/components/onvif/camera.py

@@ -505,4 +505,6 @@ class ONVIFHassCamera(Camera):
     @property
     def unique_id(self) -> Optional[str]:
         """Return a unique ID."""
+        if self._profile_index:
+            return f"{self._mac}_{self._profile_index}"
         return self._mac

homeassistant/components/samsungtv/bridge.py

@@ -206,6 +206,7 @@ class SamsungTVWSBridge(SamsungTVBridge):
                 CONF_TIMEOUT: 31,
             }
 
+            result = None
             try:
                 LOGGER.debug("Try config: %s", config)
                 with SamsungTVWS(
@@ -223,9 +224,13 @@ class SamsungTVWSBridge(SamsungTVBridge):
                 return RESULT_SUCCESS
             except WebSocketException:
                 LOGGER.debug("Working but unsupported config: %s", config)
-                return RESULT_NOT_SUPPORTED
+                result = RESULT_NOT_SUPPORTED
             except (OSError, ConnectionFailure) as err:
                 LOGGER.debug("Failing config: %s, error: %s", config, err)
+        # pylint: disable=useless-else-on-loop
+        else:
+            if result:
+                return result
 
         return RESULT_NOT_SUCCESSFUL
 

homeassistant/components/tankerkoenig/__init__.py

@@ -1,6 +1,7 @@
 """Ask tankerkoenig.de for petrol price information."""
 from datetime import timedelta
 import logging
+from math import ceil
 
 import pytankerkoenig
 import voluptuous as vol
@@ -164,27 +165,41 @@ class TankerkoenigData:
                 )
                 return False
             self.add_station(additional_station_data["station"])
+        if len(self.stations) > 10:
+            _LOGGER.warning(
+                "Found more than 10 stations to check. "
+                "This might invalidate your api-key on the long run. "
+                "Try using a smaller radius"
+            )
         return True
 
     async def fetch_data(self):
         """Get the latest data from tankerkoenig.de."""
         _LOGGER.debug("Fetching new data from tankerkoenig.de")
         station_ids = list(self.stations)
-        data = await self._hass.async_add_executor_job(
-            pytankerkoenig.getPriceList, self._api_key, station_ids
-        )
 
-        if data["ok"]:
+        prices = {}
+
+        # The API seems to only return at most 10 results, so split the list in chunks of 10
+        # and merge it together.
+        for index in range(ceil(len(station_ids) / 10)):
+            data = await self._hass.async_add_executor_job(
+                pytankerkoenig.getPriceList,
+                self._api_key,
+                station_ids[index * 10 : (index + 1) * 10],
+            )
+
             _LOGGER.debug("Received data: %s", data)
+            if not data["ok"]:
+                _LOGGER.error(
+                    "Error fetching data from tankerkoenig.de: %s", data["message"]
+                )
+                raise TankerkoenigError(data["message"])
             if "prices" not in data:
                 _LOGGER.error("Did not receive price information from tankerkoenig.de")
                 raise TankerkoenigError("No prices in data")
-        else:
+            prices.update(data["prices"])
-            _LOGGER.error(
+        return prices
-                "Error fetching data from tankerkoenig.de: %s", data["message"]
-            )
-            raise TankerkoenigError(data["message"])
-        return data["prices"]
 
     def add_station(self, station: dict):
         """Add fuel station to the entity list."""

homeassistant/components/totalconnect/alarm_control_panel.py

@@ -18,7 +18,7 @@ from homeassistant.const import (
     STATE_ALARM_TRIGGERED,
 )
 
-from . import DOMAIN as TOTALCONNECT_DOMAIN
+from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -30,7 +30,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
 
     alarms = []
 
-    client = hass.data[TOTALCONNECT_DOMAIN].client
+    client = hass.data[DOMAIN].client
 
     for location_id, location in client.locations.items():
         location_name = location.location_name
@@ -71,7 +71,7 @@ class TotalConnectAlarm(alarm.AlarmControlPanel):
 
     def update(self):
         """Return the state of the device."""
-        status = self._client.get_armed_status(self._location_id)
+        self._client.get_armed_status(self._location_id)
         attr = {
             "location_name": self._name,
             "location_id": self._location_id,
@@ -79,47 +79,36 @@ class TotalConnectAlarm(alarm.AlarmControlPanel):
             "low_battery": self._client.locations[self._location_id].low_battery,
             "cover_tampered": self._client.locations[
                 self._location_id
-            ].is_cover_tampered,
+            ].is_cover_tampered(),
             "triggered_source": None,
             "triggered_zone": None,
         }
 
-        if status in (self._client.DISARMED, self._client.DISARMED_BYPASS):
+        if self._client.locations[self._location_id].is_disarmed():
             state = STATE_ALARM_DISARMED
-        elif status in (
+        elif self._client.locations[self._location_id].is_armed_home():
-            self._client.ARMED_STAY,
-            self._client.ARMED_STAY_INSTANT,
-            self._client.ARMED_STAY_INSTANT_BYPASS,
-        ):
             state = STATE_ALARM_ARMED_HOME
-        elif status == self._client.ARMED_STAY_NIGHT:
+        elif self._client.locations[self._location_id].is_armed_night():
             state = STATE_ALARM_ARMED_NIGHT
-        elif status in (
+        elif self._client.locations[self._location_id].is_armed_away():
-            self._client.ARMED_AWAY,
-            self._client.ARMED_AWAY_BYPASS,
-            self._client.ARMED_AWAY_INSTANT,
-            self._client.ARMED_AWAY_INSTANT_BYPASS,
-        ):
             state = STATE_ALARM_ARMED_AWAY
-        elif status == self._client.ARMED_CUSTOM_BYPASS:
+        elif self._client.locations[self._location_id].is_armed_custom_bypass():
             state = STATE_ALARM_ARMED_CUSTOM_BYPASS
-        elif status == self._client.ARMING:
+        elif self._client.locations[self._location_id].is_arming():
             state = STATE_ALARM_ARMING
-        elif status == self._client.DISARMING:
+        elif self._client.locations[self._location_id].is_disarming():
             state = STATE_ALARM_DISARMING
-        elif status == self._client.ALARMING:
+        elif self._client.locations[self._location_id].is_triggered_police():
             state = STATE_ALARM_TRIGGERED
             attr["triggered_source"] = "Police/Medical"
-        elif status == self._client.ALARMING_FIRE_SMOKE:
+        elif self._client.locations[self._location_id].is_triggered_fire():
             state = STATE_ALARM_TRIGGERED
             attr["triggered_source"] = "Fire/Smoke"
-        elif status == self._client.ALARMING_CARBON_MONOXIDE:
+        elif self._client.locations[self._location_id].is_triggered_gas():
             state = STATE_ALARM_TRIGGERED
             attr["triggered_source"] = "Carbon Monoxide"
         else:
-            logging.info(
+            logging.info("Total Connect Client returned unknown status")
-                "Total Connect Client returned unknown status code: %s", status
-            )
             state = None
 
         self._state = state

homeassistant/components/totalconnect/const.py

@@ -0,0 +1,3 @@
+"""TotalConnect constants."""
+
+DOMAIN = "totalconnect"

homeassistant/components/zha/__init__.py

@@ -8,6 +8,8 @@ import voluptuous as vol
 from homeassistant import config_entries, const as ha_const
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE
+from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.typing import HomeAssistantType
 
 from . import api
 from .core import ZHAGateway
@@ -27,6 +29,7 @@ from .core.const import (
     DEFAULT_BAUDRATE,
     DEFAULT_RADIO_TYPE,
     DOMAIN,
+    SIGNAL_ADD_ENTITIES,
     RadioType,
 )
 
@@ -90,8 +93,15 @@ async def async_setup_entry(hass, config_entry):
     """
 
     zha_data = hass.data.setdefault(DATA_ZHA, {})
+    zha_data[DATA_ZHA_PLATFORM_LOADED] = {}
     config = zha_data.get(DATA_ZHA_CONFIG, {})
 
+    zha_data[DATA_ZHA_DISPATCHERS] = []
+    for component in COMPONENTS:
+        zha_data[component] = []
+        coro = hass.config_entries.async_forward_entry_setup(config_entry, component)
+        zha_data[DATA_ZHA_PLATFORM_LOADED][component] = hass.async_create_task(coro)
+
     if config.get(CONF_ENABLE_QUIRKS, True):
         # needs to be done here so that the ZHA module is finished loading
         # before zhaquirks is imported
@@ -100,22 +110,6 @@ async def async_setup_entry(hass, config_entry):
     zha_gateway = ZHAGateway(hass, config, config_entry)
     await zha_gateway.async_initialize()
 
-    zha_data[DATA_ZHA_DISPATCHERS] = []
-    zha_data[DATA_ZHA_PLATFORM_LOADED] = asyncio.Event()
-    platforms = []
-    for component in COMPONENTS:
-        platforms.append(
-            hass.async_create_task(
-                hass.config_entries.async_forward_entry_setup(config_entry, component)
-            )
-        )
-
-    async def _platforms_loaded():
-        await asyncio.gather(*platforms)
-        zha_data[DATA_ZHA_PLATFORM_LOADED].set()
-
-    hass.async_create_task(_platforms_loaded())
-
     device_registry = await hass.helpers.device_registry.async_get_registry()
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
@@ -134,7 +128,7 @@ async def async_setup_entry(hass, config_entry):
         await zha_data[DATA_ZHA_GATEWAY].async_update_device_storage()
 
     hass.bus.async_listen_once(ha_const.EVENT_HOMEASSISTANT_STOP, async_zha_shutdown)
-    hass.async_create_task(zha_gateway.async_load_devices())
+    hass.async_create_task(async_load_entities(hass, config_entry))
     return True
 
 
@@ -152,3 +146,20 @@ async def async_unload_entry(hass, config_entry):
         await hass.config_entries.async_forward_entry_unload(config_entry, component)
 
     return True
+
+
+async def async_load_entities(
+    hass: HomeAssistantType, config_entry: config_entries.ConfigEntry
+) -> None:
+    """Load entities after integration was setup."""
+    await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_prepare_entities()
+    to_setup = [
+        hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED][comp]
+        for comp in COMPONENTS
+        if hass.data[DATA_ZHA][comp]
+    ]
+    results = await asyncio.gather(*to_setup, return_exceptions=True)
+    for res in results:
+        if isinstance(res, Exception):
+            _LOGGER.warning("Couldn't setup zha platform: %s", res)
+    async_dispatcher_send(hass, SIGNAL_ADD_ENTITIES)

homeassistant/components/zha/binary_sensor.py

@@ -49,7 +49,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation binary sensor from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/core/discovery.py

@@ -8,6 +8,16 @@ from homeassistant.core import callback
 from homeassistant.helpers.typing import HomeAssistantType
 
 from . import const as zha_const, registries as zha_regs, typing as zha_typing
+from .. import (  # noqa: F401 pylint: disable=unused-import,
+    binary_sensor,
+    cover,
+    device_tracker,
+    fan,
+    light,
+    lock,
+    sensor,
+    switch,
+)
 from .channels import base
 
 _LOGGER = logging.getLogger(__name__)

homeassistant/components/zha/core/gateway.py

@@ -36,7 +36,6 @@ from .const import (
     DATA_ZHA,
     DATA_ZHA_BRIDGE_ID,
     DATA_ZHA_GATEWAY,
-    DATA_ZHA_PLATFORM_LOADED,
     DEBUG_COMP_BELLOWS,
     DEBUG_COMP_ZHA,
     DEBUG_COMP_ZIGPY,
@@ -157,34 +156,40 @@ class ZHAGateway:
         self._hass.data[DATA_ZHA][DATA_ZHA_BRIDGE_ID] = str(
             self.application_controller.ieee
         )
+        await self.async_load_devices()
         self._initialize_groups()
 
     async def async_load_devices(self) -> None:
         """Restore ZHA devices from zigpy application state."""
-        await self._hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED].wait()
+        zigpy_devices = self.application_controller.devices.values()
+        for zigpy_device in zigpy_devices:
+            self._async_get_or_create_device(zigpy_device, restored=True)
 
+    async def async_prepare_entities(self) -> None:
+        """Prepare entities by initializing device channels."""
         semaphore = asyncio.Semaphore(2)
 
-        async def _throttle(device: zha_typing.ZigpyDeviceType):
+        async def _throttle(zha_device: zha_typing.ZhaDeviceType, cached: bool):
             async with semaphore:
-                await self.async_device_restored(device)
+                await zha_device.async_initialize(from_cache=cached)
 
-        zigpy_devices = self.application_controller.devices.values()
         _LOGGER.debug("Loading battery powered devices")
         await asyncio.gather(
             *[
-                _throttle(dev)
+                _throttle(dev, cached=True)
-                for dev in zigpy_devices
+                for dev in self.devices.values()
-                if not dev.node_desc.is_mains_powered
+                if not dev.is_mains_powered
             ]
         )
-        async_dispatcher_send(self._hass, SIGNAL_ADD_ENTITIES)
 
         _LOGGER.debug("Loading mains powered devices")
         await asyncio.gather(
-            *[_throttle(dev) for dev in zigpy_devices if dev.node_desc.is_mains_powered]
+            *[
+                _throttle(dev, cached=False)
+                for dev in self.devices.values()
+                if dev.is_mains_powered
+            ]
         )
-        async_dispatcher_send(self._hass, SIGNAL_ADD_ENTITIES)
 
     def device_joined(self, device):
         """Handle device joined.

homeassistant/components/zha/cover.py

@@ -29,7 +29,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation cover from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/device_tracker.py

@@ -26,7 +26,7 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation device tracker from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/fan.py

@@ -53,7 +53,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation fan from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/light.py

@@ -52,7 +52,7 @@ _REFRESH_INTERVAL = (45, 75)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation light from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][light.DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][light.DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/lock.py

@@ -36,7 +36,7 @@ VALUE_TO_STATE = dict(enumerate(STATE_LIST))
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation Door Lock from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/sensor.py

@@ -68,7 +68,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation sensor from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/switch.py

@@ -26,7 +26,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation switch from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/const.py

@@ -1,7 +1,7 @@
 """Constants used by Home Assistant components."""
 MAJOR_VERSION = 0
 MINOR_VERSION = 107
-PATCH_VERSION = "4"
+PATCH_VERSION = "5"
 __short_version__ = f"{MAJOR_VERSION}.{MINOR_VERSION}"
 __version__ = f"{__short_version__}.{PATCH_VERSION}"
 REQUIRED_PYTHON_VER = (3, 7, 0)

requirements_all.txt

@@ -712,7 +712,7 @@ horimote==0.4.1
 httplib2==0.10.3
 
 # homeassistant.components.huawei_lte
-huawei-lte-api==1.4.10
+huawei-lte-api==1.4.11
 
 # homeassistant.components.hydrawise
 hydrawiser==0.1.1

requirements_test_all.txt

@@ -276,7 +276,7 @@ homematicip==0.10.17
 httplib2==0.10.3
 
 # homeassistant.components.huawei_lte
-huawei-lte-api==1.4.10
+huawei-lte-api==1.4.11
 
 # homeassistant.components.iaqualink
 iaqualink==0.3.1

tests/components/samsungtv/test_config_flow.py

@@ -1,5 +1,5 @@
 """Tests for Samsung TV config flow."""
-from unittest.mock import call, patch
+from unittest.mock import Mock, PropertyMock, call, patch
 
 from asynctest import mock
 import pytest
@@ -19,7 +19,7 @@ from homeassistant.components.ssdp import (
     ATTR_UPNP_MODEL_NAME,
     ATTR_UPNP_UDN,
 )
-from homeassistant.const import CONF_HOST, CONF_ID, CONF_METHOD, CONF_NAME
+from homeassistant.const import CONF_HOST, CONF_ID, CONF_METHOD, CONF_NAME, CONF_TOKEN
 
 MOCK_USER_DATA = {CONF_HOST: "fake_host", CONF_NAME: "fake_name"}
 MOCK_SSDP_DATA = {
@@ -46,6 +46,20 @@ AUTODETECT_LEGACY = {
     "host": "fake_host",
     "timeout": 31,
 }
+AUTODETECT_WEBSOCKET_PLAIN = {
+    "host": "fake_host",
+    "name": "HomeAssistant",
+    "port": 8001,
+    "timeout": 31,
+    "token": None,
+}
+AUTODETECT_WEBSOCKET_SSL = {
+    "host": "fake_host",
+    "name": "HomeAssistant",
+    "port": 8002,
+    "timeout": 31,
+    "token": None,
+}
 
 
 @pytest.fixture(name="remote")
@@ -446,20 +460,48 @@ async def test_autodetect_websocket(hass, remote, remotews):
     with patch(
         "homeassistant.components.samsungtv.bridge.Remote", side_effect=OSError("Boom"),
     ), patch("homeassistant.components.samsungtv.bridge.SamsungTVWS") as remotews:
+        enter = Mock()
+        type(enter).token = PropertyMock(return_value="123456789")
+        remote = Mock()
+        remote.__enter__ = Mock(return_value=enter)
+        remote.__exit__ = Mock(return_value=False)
+        remotews.return_value = remote
+
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": "user"}, data=MOCK_USER_DATA
         )
         assert result["type"] == "create_entry"
         assert result["data"][CONF_METHOD] == "websocket"
+        assert result["data"][CONF_TOKEN] == "123456789"
         assert remotews.call_count == 1
+        assert remotews.call_args_list == [call(**AUTODETECT_WEBSOCKET_PLAIN)]
+
+
+async def test_autodetect_websocket_ssl(hass, remote, remotews):
+    """Test for send key with autodetection of protocol."""
+    with patch(
+        "homeassistant.components.samsungtv.bridge.Remote", side_effect=OSError("Boom"),
+    ), patch(
+        "homeassistant.components.samsungtv.bridge.SamsungTVWS",
+        side_effect=[WebSocketProtocolException("Boom"), mock.DEFAULT],
+    ) as remotews:
+        enter = Mock()
+        type(enter).token = PropertyMock(return_value="123456789")
+        remote = Mock()
+        remote.__enter__ = Mock(return_value=enter)
+        remote.__exit__ = Mock(return_value=False)
+        remotews.return_value = remote
+
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": "user"}, data=MOCK_USER_DATA
+        )
+        assert result["type"] == "create_entry"
+        assert result["data"][CONF_METHOD] == "websocket"
+        assert result["data"][CONF_TOKEN] == "123456789"
+        assert remotews.call_count == 2
         assert remotews.call_args_list == [
-            call(
+            call(**AUTODETECT_WEBSOCKET_PLAIN),
-                host="fake_host",
+            call(**AUTODETECT_WEBSOCKET_SSL),
-                name="HomeAssistant",
-                port=8001,
-                timeout=31,
-                token=None,
-            )
         ]
 
 
@@ -524,18 +566,6 @@ async def test_autodetect_none(hass, remote, remotews):
         ]
         assert remotews.call_count == 2
         assert remotews.call_args_list == [
-            call(
+            call(**AUTODETECT_WEBSOCKET_PLAIN),
-                host="fake_host",
+            call(**AUTODETECT_WEBSOCKET_SSL),
-                name="HomeAssistant",
-                port=8001,
-                timeout=31,
-                token=None,
-            ),
-            call(
-                host="fake_host",
-                name="HomeAssistant",
-                port=8002,
-                timeout=31,
-                token=None,
-            ),
         ]

tests/components/samsungtv/test_media_player.py

@@ -34,8 +34,11 @@ from homeassistant.const import (
     ATTR_FRIENDLY_NAME,
     ATTR_SUPPORTED_FEATURES,
     CONF_HOST,
+    CONF_IP_ADDRESS,
+    CONF_METHOD,
     CONF_NAME,
     CONF_PORT,
+    CONF_TOKEN,
     SERVICE_MEDIA_NEXT_TRACK,
     SERVICE_MEDIA_PAUSE,
     SERVICE_MEDIA_PLAY,
@@ -51,7 +54,7 @@ from homeassistant.const import (
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 ENTITY_ID = f"{DOMAIN}.fake"
 MOCK_CONFIG = {
@@ -64,17 +67,40 @@ MOCK_CONFIG = {
         }
     ]
 }
-
 MOCK_CONFIGWS = {
     SAMSUNGTV_DOMAIN: [
         {
             CONF_HOST: "fake",
             CONF_NAME: "fake",
             CONF_PORT: 8001,
+            CONF_TOKEN: "123456789",
             CONF_ON_ACTION: [{"delay": "00:00:01"}],
         }
     ]
 }
+MOCK_CALLS_WS = {
+    "host": "fake",
+    "port": 8001,
+    "token": None,
+    "timeout": 31,
+    "name": "HomeAssistant",
+}
+
+MOCK_ENTRY_WS = {
+    CONF_IP_ADDRESS: "test",
+    CONF_HOST: "fake",
+    CONF_METHOD: "websocket",
+    CONF_NAME: "fake",
+    CONF_PORT: 8001,
+    CONF_TOKEN: "abcde",
+}
+MOCK_CALLS_ENTRY_WS = {
+    "host": "fake",
+    "name": "HomeAssistant",
+    "port": 8001,
+    "timeout": 1,
+    "token": "abcde",
+}
 
 ENTITY_ID_NOTURNON = f"{DOMAIN}.fake_noturnon"
 MOCK_CONFIG_NOTURNON = {
@@ -155,6 +181,52 @@ async def test_setup_without_turnon(hass, remote):
     assert hass.states.get(ENTITY_ID_NOTURNON)
 
 
+async def test_setup_websocket(hass, remotews, mock_now):
+    """Test setup of platform."""
+    with patch("homeassistant.components.samsungtv.bridge.SamsungTVWS") as remote_class:
+        enter = mock.Mock()
+        type(enter).token = mock.PropertyMock(return_value="987654321")
+        remote = mock.Mock()
+        remote.__enter__ = mock.Mock(return_value=enter)
+        remote.__exit__ = mock.Mock()
+        remote_class.return_value = remote
+
+        await setup_samsungtv(hass, MOCK_CONFIGWS)
+
+        assert remote_class.call_count == 1
+        assert remote_class.call_args_list == [call(**MOCK_CALLS_WS)]
+        assert hass.states.get(ENTITY_ID)
+
+
+async def test_setup_websocket_2(hass, mock_now):
+    """Test setup of platform from config entry."""
+    entity_id = f"{DOMAIN}.fake"
+
+    entry = MockConfigEntry(
+        domain=SAMSUNGTV_DOMAIN, data=MOCK_ENTRY_WS, unique_id=entity_id,
+    )
+    entry.add_to_hass(hass)
+
+    config_entries = hass.config_entries.async_entries(SAMSUNGTV_DOMAIN)
+    assert len(config_entries) == 1
+    assert entry is config_entries[0]
+
+    assert await async_setup_component(hass, SAMSUNGTV_DOMAIN, {})
+    await hass.async_block_till_done()
+
+    next_update = mock_now + timedelta(minutes=5)
+    with patch(
+        "homeassistant.components.samsungtv.bridge.SamsungTVWS"
+    ) as remote, patch("homeassistant.util.dt.utcnow", return_value=next_update):
+        async_fire_time_changed(hass, next_update)
+        await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert remote.call_count == 1
+    assert remote.call_args_list == [call(**MOCK_CALLS_ENTRY_WS)]
+
+
 async def test_update_on(hass, remote, mock_now):
     """Testing update tv on."""
     await setup_samsungtv(hass, MOCK_CONFIG)