diff --git a/homeassistant/components/smarttub/controller.py b/homeassistant/components/smarttub/controller.py index b1f9a3d4948..06b0989233c 100644 --- a/homeassistant/components/smarttub/controller.py +++ b/homeassistant/components/smarttub/controller.py @@ -92,16 +92,14 @@ class SmartTubController: return data async def _get_spa_data(self, spa): - status, pumps, lights, reminders = await asyncio.gather( - spa.get_status(), - spa.get_pumps(), - spa.get_lights(), + full_status, reminders = await asyncio.gather( + spa.get_status_full(), spa.get_reminders(), ) return { - ATTR_STATUS: status, - ATTR_PUMPS: {pump.id: pump for pump in pumps}, - ATTR_LIGHTS: {light.zone: light for light in lights}, + ATTR_STATUS: full_status, + ATTR_PUMPS: {pump.id: pump for pump in full_status.pumps}, + ATTR_LIGHTS: {light.zone: light for light in full_status.lights}, ATTR_REMINDERS: {reminder.id: reminder for reminder in reminders}, } diff --git a/tests/components/smarttub/conftest.py b/tests/components/smarttub/conftest.py index 7f23c2355bc..84566fcccc5 100644 --- a/tests/components/smarttub/conftest.py +++ b/tests/components/smarttub/conftest.py @@ -42,9 +42,9 @@ def mock_spa(): mock_spa.id = "mockspa1" mock_spa.brand = "mockbrand1" mock_spa.model = "mockmodel1" - mock_spa.get_status.return_value = smarttub.SpaState( + full_status = smarttub.SpaStateFull( mock_spa, - **{ + { "setTemperature": 39, "water": {"temperature": 38}, "heater": "ON", @@ -69,8 +69,12 @@ def mock_spa(): "uv": "OFF", "blowoutCycle": "INACTIVE", "cleanupCycle": "INACTIVE", + "lights": [], + "pumps": [], }, ) + mock_spa.get_status_full.return_value = full_status + mock_circulation_pump = create_autospec(smarttub.SpaPump, instance=True) mock_circulation_pump.id = "CP" mock_circulation_pump.spa = mock_spa @@ -89,7 +93,7 @@ def mock_spa(): mock_jet_on.state = smarttub.SpaPump.PumpState.HIGH mock_jet_on.type = smarttub.SpaPump.PumpType.JET - mock_spa.get_pumps.return_value = [mock_circulation_pump, mock_jet_off, mock_jet_on] + full_status.pumps = [mock_circulation_pump, mock_jet_off, mock_jet_on] mock_light_off = create_autospec(smarttub.SpaLight, instance=True) mock_light_off.spa = mock_spa @@ -103,7 +107,7 @@ def mock_spa(): mock_light_on.intensity = 50 mock_light_on.mode = smarttub.SpaLight.LightMode.PURPLE - mock_spa.get_lights.return_value = [mock_light_off, mock_light_on] + full_status.lights = [mock_light_off, mock_light_on] mock_filter_reminder = create_autospec(smarttub.SpaReminder, instance=True) mock_filter_reminder.id = "FILTER01" diff --git a/tests/components/smarttub/test_climate.py b/tests/components/smarttub/test_climate.py index a034a4ce17e..a9c2de4e6e2 100644 --- a/tests/components/smarttub/test_climate.py +++ b/tests/components/smarttub/test_climate.py @@ -42,7 +42,7 @@ async def test_thermostat_update(spa, setup_entry, hass): assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_HEAT - spa.get_status.return_value.heater = "OFF" + spa.get_status_full.return_value.heater = "OFF" await trigger_update(hass) state = hass.states.get(entity_id) @@ -85,11 +85,11 @@ async def test_thermostat_update(spa, setup_entry, hass): ) spa.set_heat_mode.assert_called_with(smarttub.Spa.HeatMode.ECONOMY) - spa.get_status.return_value.heat_mode = smarttub.Spa.HeatMode.ECONOMY + spa.get_status_full.return_value.heat_mode = smarttub.Spa.HeatMode.ECONOMY await trigger_update(hass) state = hass.states.get(entity_id) assert state.attributes.get(ATTR_PRESET_MODE) == PRESET_ECO - spa.get_status.side_effect = smarttub.APIError + spa.get_status_full.side_effect = smarttub.APIError await trigger_update(hass) # should not fail diff --git a/tests/components/smarttub/test_light.py b/tests/components/smarttub/test_light.py index fe178278bee..9f128cc279c 100644 --- a/tests/components/smarttub/test_light.py +++ b/tests/components/smarttub/test_light.py @@ -32,9 +32,8 @@ async def test_light( assert state is not None assert state.state == light_state - light: SpaLight = next( - light for light in await spa.get_lights() if light.zone == light_zone - ) + status = await spa.get_status_full() + light: SpaLight = next(light for light in status.lights if light.zone == light_zone) await hass.services.async_call( "light", diff --git a/tests/components/smarttub/test_switch.py b/tests/components/smarttub/test_switch.py index 81b53604065..f97877f6cde 100644 --- a/tests/components/smarttub/test_switch.py +++ b/tests/components/smarttub/test_switch.py @@ -16,7 +16,8 @@ from homeassistant.const import STATE_OFF, STATE_ON async def test_pumps(spa, setup_entry, hass, pump_id, pump_state, entity_suffix): """Test pump entities.""" - pump = next(pump for pump in await spa.get_pumps() if pump.id == pump_id) + status = await spa.get_status_full() + pump = next(pump for pump in status.pumps if pump.id == pump_id) entity_id = f"switch.{spa.brand}_{spa.model}_{entity_suffix}" state = hass.states.get(entity_id)