Fix config-flow tests

home-assistant · Feb 16, 2022 · 65a8d9c · 65a8d9c
1 parent a188865
commit 65a8d9c
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/tests/components/samsungtv/test_config_flow.py b/tests/components/samsungtv/test_config_flow.py
@@ -1,6 +1,6 @@
 """Tests for Samsung TV config flow."""
 import socket
-from unittest.mock import Mock, PropertyMock, call, patch
+from unittest.mock import Mock, call, patch
 
 from samsungctl.exceptions import AccessDenied, UnhandledResponse
 from samsungtvws.exceptions import ConnectionFailure, HttpApiError
@@ -828,10 +828,8 @@ async def test_autodetect_websocket(hass: HomeAssistant, remote: Mock, remotews:
     ), 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.__enter__ = Mock(return_value=remote)
         remote.__exit__ = Mock(return_value=False)
         remote.rest_device_info.return_value = {
             "id": "uuid:be9554b9-c9fb-41f4-8920-22da015376a4",
@@ -845,6 +843,7 @@ async def test_autodetect_websocket(hass: HomeAssistant, remote: Mock, remotews:
                 "type": "Samsung SmartTV",
             },
         }
+        remote.token = "123456789"
         remotews.return_value = remote
 
         result = await hass.config_entries.flow.async_init(
@@ -878,10 +877,8 @@ async def test_websocket_no_mac(hass: HomeAssistant, remote: Mock, remotews: Moc
     ) as remotews, patch(
         "getmac.get_mac_address", return_value="gg:hh:ii:ll:mm:nn"
     ):
-        enter = Mock()
-        type(enter).token = PropertyMock(return_value="123456789")
         remote = Mock()
-        remote.__enter__ = Mock(return_value=enter)
+        remote.__enter__ = Mock(return_value=remote)
         remote.__exit__ = Mock(return_value=False)
         remote.rest_device_info.return_value = {
             "id": "uuid:be9554b9-c9fb-41f4-8920-22da015376a4",
@@ -893,6 +890,7 @@ async def test_websocket_no_mac(hass: HomeAssistant, remote: Mock, remotews: Moc
                 "type": "Samsung SmartTV",
             },
         }
+        remote.token = "123456789"
         remotews.return_value = remote
 
         result = await hass.config_entries.flow.async_init(