update tests

Author: jeeftor

tests/components/weatherflow_cloud/conftest.py

@@ -1,14 +1,16 @@
 """Common fixtures for the WeatherflowCloud tests."""
 
 from collections.abc import Generator
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 from aiohttp import ClientResponseError
 import pytest
+from weatherflow4py.api import WeatherFlowRestAPI
 from weatherflow4py.models.rest.forecast import WeatherDataForecastREST
 from weatherflow4py.models.rest.observation import ObservationStationREST
 from weatherflow4py.models.rest.stations import StationsResponseREST
 from weatherflow4py.models.rest.unified import WeatherFlowDataREST
+from weatherflow4py.ws import WeatherFlowWebsocketAPI
 
 from homeassistant.components.weatherflow_cloud.const import DOMAIN
 from homeassistant.const import CONF_API_TOKEN
@@ -81,35 +83,83 @@ async def mock_config_entry() -> MockConfigEntry:
 
 
 @pytest.fixture
-def mock_api():
-    """Fixture for Mock WeatherFlowRestAPI."""
-    get_stations_response_data = StationsResponseREST.from_json(
-        load_fixture("stations.json", DOMAIN)
-    )
-    get_forecast_response_data = WeatherDataForecastREST.from_json(
-        load_fixture("forecast.json", DOMAIN)
-    )
-    get_observation_response_data = ObservationStationREST.from_json(
-        load_fixture("station_observation.json", DOMAIN)
-    )
+def mock_rest_api():
+    """Mock rest api."""
+    fixtures = {
+        "stations": StationsResponseREST.from_json(
+            load_fixture("stations.json", DOMAIN)
+        ),
+        "forecast": WeatherDataForecastREST.from_json(
+            load_fixture("forecast.json", DOMAIN)
+        ),
+        "observation": ObservationStationREST.from_json(
+            load_fixture("station_observation.json", DOMAIN)
+        ),
+    }
+
+    # Create device_station_map
+    device_station_map = {
+        device.device_id: station.station_id
+        for station in fixtures["stations"].stations
+        for device in station.devices
+    }
 
+    # Prepare mock data
     data = {
         24432: WeatherFlowDataREST(
-            weather=get_forecast_response_data,
-            observation=get_observation_response_data,
-            station=get_stations_response_data.stations[0],
+            weather=fixtures["forecast"],
+            observation=fixtures["observation"],
+            station=fixtures["stations"].stations[0],
             device_observations=None,
         )
     }
 
-    with patch(
-        "homeassistant.components.weatherflow_cloud.coordinator.WeatherFlowRestAPI",
-        autospec=True,
-    ) as mock_api_class:
-        # Create an instance of AsyncMock for the API
-        mock_api = AsyncMock()
-        mock_api.get_all_data.return_value = data
-        # Patch the class to return our mock_api instance
-        mock_api_class.return_value = mock_api
-
+    mock_api = AsyncMock(spec=WeatherFlowRestAPI)
+    mock_api.get_all_data.return_value = data
+    mock_api.async_get_stations.return_value = fixtures["stations"]
+    mock_api.device_station_map = device_station_map
+    mock_api.api_token = MOCK_API_TOKEN
+
+    # Apply patches
+    with (
+        patch(
+            "homeassistant.components.weatherflow_cloud.WeatherFlowRestAPI",
+            return_value=mock_api,
+        ) as _,
+        patch(
+            "homeassistant.components.weatherflow_cloud.coordinator.WeatherFlowRestAPI",
+            return_value=mock_api,
+        ) as _,
+    ):
         yield mock_api
+
+
+@pytest.fixture
+async def mock_websocket_api():
+    """Mock WeatherFlowWebsocketAPI."""
+    mock_websocket = AsyncMock()
+    mock_websocket.send = AsyncMock()
+    mock_websocket.recv = AsyncMock()
+
+    mock_ws_instance = AsyncMock(spec=WeatherFlowWebsocketAPI)
+    mock_ws_instance.connect = AsyncMock()
+    mock_ws_instance.send_message = AsyncMock()
+    mock_ws_instance.register_observation_callback = MagicMock()
+    mock_ws_instance.register_wind_callback = MagicMock()
+    mock_ws_instance.websocket = mock_websocket
+
+    with (
+        patch(
+            "homeassistant.components.weatherflow_cloud.coordinator.WeatherFlowWebsocketAPI",
+            return_value=mock_ws_instance,
+        ),
+        patch(
+            "homeassistant.components.weatherflow_cloud.WeatherFlowWebsocketAPI",
+            return_value=mock_ws_instance,
+        ),
+        patch(
+            "weatherflow4py.ws.WeatherFlowWebsocketAPI", return_value=mock_ws_instance
+        ),
+    ):
+        # mock_connect.return_value = mock_websocket
+        yield mock_ws_instance