diff --git a/src/homematicip/connection/websocket_handler.py b/src/homematicip/connection/websocket_handler.py index 0d36a7d..cdad84e 100644 --- a/src/homematicip/connection/websocket_handler.py +++ b/src/homematicip/connection/websocket_handler.py @@ -25,12 +25,12 @@ async def listen(self, context: ConnectionContext, connection_state_callback: Ca async for message in websocket: try: if connection_state_callback is not None: - connection_state_callback(True) + await connection_state_callback(True) yield message except websockets.ConnectionClosed: if connection_state_callback is not None: - connection_state_callback(False) + await connection_state_callback(False) LOGGER.warn( f"Got ConnectionClosed Exception. Should reconnect is set to {reconnect_on_error}" ) @@ -40,4 +40,4 @@ async def listen(self, context: ConnectionContext, connection_state_callback: Ca continue # If we reach this point, the connection is closed. - connection_state_callback(False) + await connection_state_callback(False) diff --git a/src/homematicip/runner.py b/src/homematicip/runner.py index 33a3b6a..bf98e70 100644 --- a/src/homematicip/runner.py +++ b/src/homematicip/runner.py @@ -140,7 +140,7 @@ async def async_refresh_model(self): async def _async_start_listening_for_updates(self, context: ConnectionContext): self.websocket_handler = WebSocketHandler() hmip_event_handler = HmipEventHandler(event_manager=self.event_manager, model=self.model) - async for message in self.websocket_handler.listen(context, self._set_websocket_connected_state, False): + async for message in self.websocket_handler.listen(context, self._async_set_websocket_connected_state, False): try: await hmip_event_handler.process_event_async(json.loads(message)) @@ -157,7 +157,7 @@ def home_is_connected(self) -> bool: """Return if the home is connected.""" return self.model.home.connected if self.model is not None else False - def _set_websocket_connected_state(self, value): + async def _async_set_websocket_connected_state(self, value): """Set the websocket connected state and publish the event. :param value: The new state of the websocket connection.""" @@ -165,16 +165,16 @@ def _set_websocket_connected_state(self, value): self._websocket_connected = value if state_changed: - self._publish_websocket_connected(value) + await self._async_publish_websocket_connected(value) - def _publish_websocket_connected(self, connected_value: bool): + async def _async_publish_websocket_connected(self, connected_value: bool): """Publish the websocket connected event. :param connected_value: The new state of the websocket connection.""" if connected_value: - self.event_manager.async_publish(ModelUpdateEvent.HOME_CONNECTED, connected_value) + await self.event_manager.async_publish(ModelUpdateEvent.HOME_CONNECTED, connected_value) else: - self.event_manager.async_publish(ModelUpdateEvent.HOME_DISCONNECTED, connected_value) + await self.event_manager.async_publish(ModelUpdateEvent.HOME_DISCONNECTED, connected_value) @property def rest_connection(self): diff --git a/tests/test_runner.py b/tests/test_runner.py index 43eac72..fa2cf08 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -52,34 +52,37 @@ async def test_runner_async_run_home(mocker, sample_data_complete, config, conne assert runner.model is not None -def test_runner_websocket_connected_callable(): +@pytest.mark.asyncio +async def test_runner_websocket_connected_callable(): runner = Runner() assert runner.websocket_connected is False - runner._set_websocket_connected_state(True) + await runner._async_set_websocket_connected_state(True) assert runner.websocket_connected is True assert runner._websocket_connected is True -def test_runner_websocket_connected_event_is_raised(mocker): +@pytest.mark.asyncio +async def test_runner_websocket_connected_event_is_raised(mocker): """Test that the HOME_CONNECTED event is raised when the websocket connection is established.""" - mock_publish = mocker.Mock() + mock_publish = AsyncMock() runner = Runner() runner.event_manager.async_publish = mock_publish runner._websocket_connected = False - runner._set_websocket_connected_state(True) + await runner._async_set_websocket_connected_state(True) mock_publish.assert_called_with(ModelUpdateEvent.HOME_CONNECTED, True) -def test_runner_websocket_disconnected_event_is_raised(mocker): +@pytest.mark.asyncio +async def test_runner_websocket_disconnected_event_is_raised(mocker): """Test that the HOME_DISCONNECTED event is raised when the websocket connection is lost.""" - mock_publish = mocker.Mock() + mock_publish = AsyncMock() runner = Runner() runner.event_manager.async_publish = mock_publish runner._websocket_connected = True - runner._set_websocket_connected_state(False) + await runner._async_set_websocket_connected_state(False) mock_publish.assert_called_with(ModelUpdateEvent.HOME_DISCONNECTED, False)