Improve balboa tests (#61691)

.coveragerc

@@ -94,7 +94,6 @@ omit =
     homeassistant/components/azure_service_bus/*
     homeassistant/components/baidu/tts.py
     homeassistant/components/balboa/__init__.py
-    homeassistant/components/balboa/entity.py
     homeassistant/components/beewi_smartclim/sensor.py
     homeassistant/components/bbb_gpio/*
     homeassistant/components/bbox/device_tracker.py

homeassistant/components/balboa/climate.py

@@ -29,6 +29,8 @@ from homeassistant.const import (
 from .const import CLIMATE, CLIMATE_SUPPORTED_FANSTATES, CLIMATE_SUPPORTED_MODES, DOMAIN
 from .entity import BalboaEntity
 
+SET_TEMPERATURE_WAIT = 1
+
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the spa climate device."""
@@ -124,10 +126,10 @@ class BalboaSpaClimate(BalboaEntity, ClimateEntity):
         newtemp = kwargs[ATTR_TEMPERATURE]
         if newtemp > self._client.tmax[self._client.TEMPRANGE_LOW][scale]:
             await self._client.change_temprange(self._client.TEMPRANGE_HIGH)
-            await asyncio.sleep(1)
+            await asyncio.sleep(SET_TEMPERATURE_WAIT)
         if newtemp < self._client.tmin[self._client.TEMPRANGE_HIGH][scale]:
             await self._client.change_temprange(self._client.TEMPRANGE_LOW)
-            await asyncio.sleep(1)
+            await asyncio.sleep(SET_TEMPERATURE_WAIT)
         await self._client.send_temp_change(newtemp)
 
     async def async_set_preset_mode(self, preset_mode) -> None:

tests/components/balboa/__init__.py

@@ -1,7 +1,4 @@
 """Test the Balboa Spa Client integration."""
-import asyncio
-from unittest.mock import patch
-
 from homeassistant.components.balboa.const import DOMAIN as BALBOA_DOMAIN
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
@@ -22,146 +19,7 @@ async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
     )
     config_entry.add_to_hass(hass)
 
-    with patch(
+    await hass.config_entries.async_setup(config_entry.entry_id)
-        "homeassistant.components.balboa.BalboaSpaWifi",
+    await hass.async_block_till_done()
-        new=BalboaMock,
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
 
     return config_entry
-
-
-async def init_integration_mocked(hass: HomeAssistant) -> MockConfigEntry:
-    """Mock integration setup."""
-    config_entry = MockConfigEntry(
-        domain=BALBOA_DOMAIN,
-        data={
-            CONF_HOST: TEST_HOST,
-        },
-    )
-    config_entry.add_to_hass(hass)
-
-    with patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.connect",
-        new=BalboaMock.connect,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.listen_until_configured",
-        new=BalboaMock.listen_until_configured,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.listen",
-        new=BalboaMock.listen,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.check_connection_status",
-        new=BalboaMock.check_connection_status,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.send_panel_req",
-        new=BalboaMock.send_panel_req,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.send_mod_ident_req",
-        new=BalboaMock.send_mod_ident_req,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.spa_configured",
-        new=BalboaMock.spa_configured,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_model_name",
-        new=BalboaMock.get_model_name,
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
-
-    return config_entry
-
-
-class BalboaMock:
-    """Mock pybalboa library."""
-
-    def __init__(self, hostname, port=BALBOA_DEFAULT_PORT):
-        """Mock init."""
-        self.host = hostname
-        self.port = port
-        self.connected = False
-        self.new_data_cb = None
-        self.lastupd = 0
-        self.connected = False
-        self.fake_action = False
-
-    async def connect(self):
-        """Connect to the spa."""
-        self.connected = True
-        return True
-
-    async def broken_connect(self):
-        """Connect to the spa."""
-        self.connected = False
-        return False
-
-    async def disconnect(self):
-        """Stop talking to the spa."""
-        self.connected = False
-
-    async def send_panel_req(self, arg_ba, arg_bb):
-        """Send a panel request, 2 bytes of data."""
-        self.fake_action = False
-        return
-
-    async def send_mod_ident_req(self):
-        """Ask for the module identification."""
-        self.fake_action = False
-        return
-
-    @staticmethod
-    def get_macaddr():
-        """Return the macaddr of the spa wifi."""
-        return "ef:ef:ef:c0:ff:ee"
-
-    def get_model_name(self):
-        """Return the model name."""
-        self.fake_action = False
-        return "FakeSpa"
-
-    @staticmethod
-    def get_ssid():
-        """Return the software version."""
-        return "V0.0"
-
-    @staticmethod
-    async def set_time(new_time, timescale=None):
-        """Set time on spa to new_time with optional timescale."""
-        return
-
-    async def listen(self):
-        """Listen to the spa babble forever."""
-        while True:
-            if not self.connected:
-                # sleep and hope the checker fixes us
-                await asyncio.sleep(5)
-                continue
-
-            # fake it
-            await asyncio.sleep(5)
-
-    async def check_connection_status(self):
-        """Set this up to periodically check the spa connection and fix."""
-        self.fake_action = False
-        while True:
-            # fake it
-            await asyncio.sleep(15)
-
-    async def spa_configured(self):
-        """Check if the spa has been configured."""
-        self.fake_action = False
-        return
-
-    async def int_new_data_cb(self):
-        """Call false internal data callback."""
-
-        if self.new_data_cb is None:
-            return
-        await self.new_data_cb()  # pylint: disable=not-callable
-
-    async def listen_until_configured(self, maxiter=20):
-        """Listen to the spa babble until we are configured."""
-        if not self.connected:
-            return False
-        return True

tests/components/balboa/conftest.py

@@ -0,0 +1,94 @@
+"""Provide common fixtures."""
+from __future__ import annotations
+
+from collections.abc import Generator
+import time
+from unittest.mock import MagicMock, patch
+
+from pybalboa.balboa import text_heatmode
+import pytest
+
+
+@pytest.fixture(name="client")
+def client_fixture() -> Generator[None, MagicMock, None]:
+    """Mock balboa."""
+    with patch(
+        "homeassistant.components.balboa.BalboaSpaWifi", autospec=True
+    ) as mock_balboa:
+        # common attributes
+        client = mock_balboa.return_value
+        client.connected = True
+        client.lastupd = time.time()
+        client.new_data_cb = None
+        client.connect.return_value = True
+        client.get_macaddr.return_value = "ef:ef:ef:c0:ff:ee"
+        client.get_model_name.return_value = "FakeSpa"
+        client.get_ssid.return_value = "V0.0"
+
+        # constants should preferebly be moved in the library
+        # to be class attributes or further refactored
+        client.TSCALE_C = 1
+        client.TSCALE_F = 0
+        client.HEATMODE_READY = 0
+        client.HEATMODE_REST = 1
+        client.HEATMODE_RNR = 2
+        client.TIMESCALE_12H = 0
+        client.TIMESCALE_24H = 1
+        client.PUMP_OFF = 0
+        client.PUMP_LOW = 1
+        client.PUMP_HIGH = 2
+        client.TEMPRANGE_LOW = 0
+        client.TEMPRANGE_HIGH = 1
+        client.tmin = [
+            [50.0, 10.0],
+            [80.0, 26.0],
+        ]
+        client.tmax = [
+            [80.0, 26.0],
+            [104.0, 40.0],
+        ]
+        client.BLOWER_OFF = 0
+        client.BLOWER_LOW = 1
+        client.BLOWER_MEDIUM = 2
+        client.BLOWER_HIGH = 3
+        client.FILTER_OFF = 0
+        client.FILTER_1 = 1
+        client.FILTER_2 = 2
+        client.FILTER_1_2 = 3
+        client.OFF = 0
+        client.ON = 1
+        client.HEATSTATE_IDLE = 0
+        client.HEATSTATE_HEATING = 1
+        client.HEATSTATE_HEAT_WAITING = 2
+        client.VOLTAGE_240 = 240
+        client.VOLTAGE_UNKNOWN = 0
+        client.HEATERTYPE_STANDARD = "Standard"
+        client.HEATERTYPE_UNKNOWN = "Unknown"
+
+        # Climate attributes
+        client.heatmode = 0
+        client.get_heatmode_stringlist.return_value = text_heatmode
+        client.get_tempscale.return_value = client.TSCALE_F
+        client.have_blower.return_value = False
+
+        # Climate methods
+        client.get_heatstate.return_value = 0
+        client.get_blower.return_value = 0
+        client.get_curtemp.return_value = 20.0
+        client.get_settemp.return_value = 20.0
+
+        def get_heatmode(text=False):
+            """Ask for the current heatmode."""
+            if text:
+                return text_heatmode[client.heatmode]
+            return client.heatmode
+
+        client.get_heatmode.side_effect = get_heatmode
+        yield client
+
+
+@pytest.fixture(autouse=True)
+def set_temperature_wait():
+    """Mock set temperature wait time."""
+    with patch("homeassistant.components.balboa.climate.SET_TEMPERATURE_WAIT", new=0):
+        yield

tests/components/balboa/test_binary_sensor.py

@@ -1,14 +1,10 @@
 """Tests of the climate entity of the balboa integration."""
+from unittest.mock import MagicMock
 
-from unittest.mock import patch
-
-from homeassistant.components.balboa.const import DOMAIN as BALBOA_DOMAIN, SIGNAL_UPDATE
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.setup import async_setup_component
 
-from . import init_integration_mocked
+from . import init_integration
 
 ENTITY_BINARY_SENSOR = "binary_sensor.fakespa_"
 
@@ -20,57 +16,41 @@ FILTER_MAP = [
 ]
 
 
-async def test_filters(hass: HomeAssistant):
+async def test_filters(hass: HomeAssistant, client: MagicMock) -> None:
     """Test spa filters."""
 
-    config_entry = await _setup_binary_sensor_test(hass)
+    config_entry = await init_integration(hass)
 
     for filter_mode in range(4):
         for spa_filter in range(1, 3):
-            state = await _patch_filter(hass, config_entry, filter_mode, spa_filter)
+            state = await _patch_filter(
+                hass, config_entry, filter_mode, spa_filter, client
+            )
             assert state.state == FILTER_MAP[filter_mode][spa_filter - 1]
 
 
-async def test_circ_pump(hass: HomeAssistant):
+async def test_circ_pump(hass: HomeAssistant, client: MagicMock) -> None:
     """Test spa circ pump."""
-    with patch(
+    client.have_circ_pump.return_value = (True,)
-        "homeassistant.components.balboa.BalboaSpaWifi.have_circ_pump",
+    config_entry = await init_integration(hass)
-        return_value=True,
-    ):
-        config_entry = await _setup_binary_sensor_test(hass)
 
-    state = await _patch_circ_pump(hass, config_entry, True)
+    state = await _patch_circ_pump(hass, config_entry, True, client)
     assert state.state == STATE_ON
-    state = await _patch_circ_pump(hass, config_entry, False)
+    state = await _patch_circ_pump(hass, config_entry, False, client)
     assert state.state == STATE_OFF
 
 
-async def _patch_circ_pump(hass, config_entry, pump_state):
+async def _patch_circ_pump(hass, config_entry, pump_state, client):
     """Patch the circ pump state."""
-    with patch(
+    client.get_circ_pump.return_value = pump_state
-        "homeassistant.components.balboa.BalboaSpaWifi.get_circ_pump",
+    await client.new_data_cb()
-        return_value=pump_state,
-    ):
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
-        return hass.states.get(f"{ENTITY_BINARY_SENSOR}circ_pump")
-
-
-async def _patch_filter(hass, config_entry, filter_mode, num):
-    """Patch the filter state."""
-    with patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_filtermode",
-        return_value=filter_mode,
-    ):
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
-        return hass.states.get(f"{ENTITY_BINARY_SENSOR}filter{num}")
-
-
-async def _setup_binary_sensor_test(hass):
-    """Prepare the test."""
-    config_entry = await init_integration_mocked(hass)
-    await async_setup_component(hass, BALBOA_DOMAIN, config_entry)
     await hass.async_block_till_done()
+    return hass.states.get(f"{ENTITY_BINARY_SENSOR}circ_pump")
 
-    return config_entry
+
+async def _patch_filter(hass, config_entry, filter_mode, num, client):
+    """Patch the filter state."""
+    client.get_filtermode.return_value = filter_mode
+    await client.new_data_cb()
+    await hass.async_block_till_done()
+    return hass.states.get(f"{ENTITY_BINARY_SENSOR}filter{num}")

tests/components/balboa/test_climate.py

@@ -1,10 +1,10 @@
 """Tests of the climate entity of the balboa integration."""
+from __future__ import annotations
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
-from homeassistant.components.balboa.const import DOMAIN as BALBOA_DOMAIN, SIGNAL_UPDATE
 from homeassistant.components.climate.const import (
     ATTR_FAN_MODE,
     ATTR_HVAC_ACTION,
@@ -28,10 +28,8 @@ from homeassistant.components.climate.const import (
 )
 from homeassistant.const import ATTR_TEMPERATURE, TEMP_FAHRENHEIT
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.setup import async_setup_component
 
-from . import init_integration_mocked
+from . import init_integration
 
 from tests.components.climate import common
 
@@ -51,13 +49,13 @@ HVAC_SETTINGS = [
 ENTITY_CLIMATE = "climate.fakespa_climate"
 
 
-async def test_spa_defaults(hass: HomeAssistant):
+async def test_spa_defaults(hass: HomeAssistant, client: MagicMock) -> None:
     """Test supported features flags."""
-
+    await init_integration(hass)
-    await _setup_climate_test(hass)
 
     state = hass.states.get(ENTITY_CLIMATE)
 
+    assert state
     assert (
         state.attributes["supported_features"]
         == SUPPORT_TARGET_TEMPERATURE | SUPPORT_PRESET_MODE
@@ -69,16 +67,15 @@ async def test_spa_defaults(hass: HomeAssistant):
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_IDLE
 
 
-async def test_spa_defaults_fake_tscale(hass: HomeAssistant):
+async def test_spa_defaults_fake_tscale(hass: HomeAssistant, client: MagicMock) -> None:
     """Test supported features flags."""
+    client.get_tempscale.return_value = 1
 
-    with patch(
+    await init_integration(hass)
-        "homeassistant.components.balboa.BalboaSpaWifi.get_tempscale", return_value=1
-    ):
-        await _setup_climate_test(hass)
 
     state = hass.states.get(ENTITY_CLIMATE)
 
+    assert state
     assert (
         state.attributes["supported_features"]
         == SUPPORT_TARGET_TEMPERATURE | SUPPORT_PRESET_MODE
@@ -90,20 +87,19 @@ async def test_spa_defaults_fake_tscale(hass: HomeAssistant):
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_IDLE
 
 
-async def test_spa_with_blower(hass: HomeAssistant):
+async def test_spa_with_blower(hass: HomeAssistant, client: MagicMock) -> None:
     """Test supported features flags."""
+    client.have_blower.return_value = True
 
-    with patch(
+    config_entry = await init_integration(hass)
-        "homeassistant.components.balboa.BalboaSpaWifi.have_blower", return_value=True
-    ):
-        config_entry = await _setup_climate_test(hass)
 
     # force a refresh
-    async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
+    await client.new_data_cb()
     await hass.async_block_till_done()
 
     state = hass.states.get(ENTITY_CLIMATE)
 
+    assert state
     assert (
         state.attributes["supported_features"]
         == SUPPORT_TARGET_TEMPERATURE | SUPPORT_PRESET_MODE | SUPPORT_FAN_MODE
@@ -111,161 +107,144 @@ async def test_spa_with_blower(hass: HomeAssistant):
 
     for fan_state in range(4):
         # set blower
-        state = await _patch_blower(hass, config_entry, fan_state)
+        state = await _patch_blower(hass, config_entry, fan_state, client)
+        assert state
         assert state.attributes[ATTR_FAN_MODE] == FAN_SETTINGS[fan_state]
 
     # test the nonsense checks
-    for fan_state in (None, 70):
+    for fan_state in (None, 70):  # type: ignore[assignment]
-        state = await _patch_blower(hass, config_entry, fan_state)
+        state = await _patch_blower(hass, config_entry, fan_state, client)
+        assert state
         assert state.attributes[ATTR_FAN_MODE] == FAN_OFF
 
 
-async def test_spa_temperature(hass: HomeAssistant):
+async def test_spa_temperature(hass: HomeAssistant, client: MagicMock) -> None:
     """Test spa temperature settings."""
 
-    config_entry = await _setup_climate_test(hass)
+    config_entry = await init_integration(hass)
 
     # flip the spa into F
     # set temp to a valid number
-    state = await _patch_spa_settemp(hass, config_entry, 0, 100.0)
+    state = await _patch_spa_settemp(hass, config_entry, 0, 100.0, client)
+    assert state
     assert state.attributes.get(ATTR_TEMPERATURE) == 38.0
 
 
-async def test_spa_temperature_unit(hass: HomeAssistant):
+async def test_spa_temperature_unit(hass: HomeAssistant, client: MagicMock) -> None:
     """Test temperature unit conversions."""
 
     with patch.object(hass.config.units, "temperature_unit", TEMP_FAHRENHEIT):
-        config_entry = await _setup_climate_test(hass)
+        config_entry = await init_integration(hass)
 
-        state = await _patch_spa_settemp(hass, config_entry, 0, 15.4)
+        state = await _patch_spa_settemp(hass, config_entry, 0, 15.4, client)
+        assert state
         assert state.attributes.get(ATTR_TEMPERATURE) == 15.0
 
 
-async def test_spa_hvac_modes(hass: HomeAssistant):
+async def test_spa_hvac_modes(hass: HomeAssistant, client: MagicMock) -> None:
     """Test hvac modes."""
 
-    config_entry = await _setup_climate_test(hass)
+    config_entry = await init_integration(hass)
 
     # try out the different heat modes
     for heat_mode in range(2):
-        state = await _patch_spa_heatmode(hass, config_entry, heat_mode)
+        state = await _patch_spa_heatmode(hass, config_entry, heat_mode, client)
+        assert state
         modes = state.attributes.get(ATTR_HVAC_MODES)
         assert [HVAC_MODE_AUTO, HVAC_MODE_HEAT, HVAC_MODE_OFF] == modes
         assert state.state == HVAC_SETTINGS[heat_mode]
 
     with pytest.raises(ValueError):
-        await _patch_spa_heatmode(hass, config_entry, 2)
+        await _patch_spa_heatmode(hass, config_entry, 2, client)
 
 
-async def test_spa_hvac_action(hass: HomeAssistant):
+async def test_spa_hvac_action(hass: HomeAssistant, client: MagicMock) -> None:
     """Test setting of the HVAC action."""
 
-    config_entry = await _setup_climate_test(hass)
+    config_entry = await init_integration(hass)
 
     # try out the different heat states
-    state = await _patch_spa_heatstate(hass, config_entry, 1)
+    state = await _patch_spa_heatstate(hass, config_entry, 1, client)
+    assert state
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_HEAT
 
-    state = await _patch_spa_heatstate(hass, config_entry, 0)
+    state = await _patch_spa_heatstate(hass, config_entry, 0, client)
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_IDLE
 
 
-async def test_spa_preset_modes(hass: HomeAssistant):
+async def test_spa_preset_modes(hass: HomeAssistant, client: MagicMock) -> None:
     """Test the various preset modes."""
 
-    config_entry = await _setup_climate_test(hass)
+    await init_integration(hass)
 
     state = hass.states.get(ENTITY_CLIMATE)
+    assert state
     modes = state.attributes.get(ATTR_PRESET_MODES)
     assert ["Ready", "Rest", "Ready in Rest"] == modes
 
     # Put it in Ready and Rest
     modelist = ["Ready", "Rest"]
     for mode in modelist:
-        with patch(
+        client.heatmode = modelist.index(mode)
-            "homeassistant.components.balboa.BalboaSpaWifi.get_heatmode",
+        await common.async_set_preset_mode(hass, mode, ENTITY_CLIMATE)
-            return_value=modelist.index(mode),
+        await client.new_data_cb()
-        ):
+        await hass.async_block_till_done()
-            await common.async_set_preset_mode(hass, mode, ENTITY_CLIMATE)
-            async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-            await hass.async_block_till_done()
 
         state = hass.states.get(ENTITY_CLIMATE)
-        assert state.attributes[ATTR_PRESET_MODE] == modelist.index(mode)
+        assert state
+        assert state.attributes[ATTR_PRESET_MODE] == mode
 
     # put it in RNR and test assertion
-    with patch(
+    client.heatmode = 2
-        "homeassistant.components.balboa.BalboaSpaWifi.get_heatmode",
+
-        return_value=2,
+    with pytest.raises(ValueError):
-    ), pytest.raises(ValueError):
         await common.async_set_preset_mode(hass, 2, ENTITY_CLIMATE)
 
 
 # Helpers
-async def _patch_blower(hass, config_entry, fan_state):
+async def _patch_blower(hass, config_entry, fan_state, client):
     """Patch the blower state."""
-    with patch(
+    client.get_blower.return_value = fan_state
-        "homeassistant.components.balboa.BalboaSpaWifi.get_blower",
-        return_value=fan_state,
-    ):
-        if fan_state is not None and fan_state <= len(FAN_SETTINGS):
-            await common.async_set_fan_mode(hass, FAN_SETTINGS[fan_state])
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
 
-    return hass.states.get(ENTITY_CLIMATE)
+    if fan_state is not None and fan_state <= len(FAN_SETTINGS):
-
+        await common.async_set_fan_mode(hass, FAN_SETTINGS[fan_state])
-
+    await client.new_data_cb()
-async def _patch_spa_settemp(hass, config_entry, tscale, settemp):
-    """Patch the settemp."""
-    with patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_tempscale",
-        return_value=tscale,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_settemp",
-        return_value=settemp,
-    ):
-        await common.async_set_temperature(
-            hass, temperature=settemp, entity_id=ENTITY_CLIMATE
-        )
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
-
-    return hass.states.get(ENTITY_CLIMATE)
-
-
-async def _patch_spa_heatmode(hass, config_entry, heat_mode):
-    """Patch the heatmode."""
-    with patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_heatmode",
-        return_value=heat_mode,
-    ):
-        await common.async_set_hvac_mode(hass, HVAC_SETTINGS[heat_mode], ENTITY_CLIMATE)
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
-
-    return hass.states.get(ENTITY_CLIMATE)
-
-
-async def _patch_spa_heatstate(hass, config_entry, heat_state):
-    """Patch the heatmode."""
-    with patch(
-        "homeassistant.components.balboa.BalboaSpaWifi.get_heatstate",
-        return_value=heat_state,
-    ):
-        await common.async_set_hvac_mode(
-            hass, HVAC_SETTINGS[heat_state], ENTITY_CLIMATE
-        )
-        async_dispatcher_send(hass, SIGNAL_UPDATE.format(config_entry.entry_id))
-        await hass.async_block_till_done()
-
-    return hass.states.get(ENTITY_CLIMATE)
-
-
-async def _setup_climate_test(hass):
-    """Prepare the test."""
-    config_entry = await init_integration_mocked(hass)
-    await async_setup_component(hass, BALBOA_DOMAIN, config_entry)
     await hass.async_block_till_done()
 
-    return config_entry
+    return hass.states.get(ENTITY_CLIMATE)
+
+
+async def _patch_spa_settemp(hass, config_entry, tscale, settemp, client):
+    """Patch the settemp."""
+    client.get_tempscale.return_value = tscale
+    client.get_settemp.return_value = settemp
+
+    await common.async_set_temperature(
+        hass, temperature=settemp, entity_id=ENTITY_CLIMATE
+    )
+    await client.new_data_cb()
+    await hass.async_block_till_done()
+
+    return hass.states.get(ENTITY_CLIMATE)
+
+
+async def _patch_spa_heatmode(hass, config_entry, heat_mode, client):
+    """Patch the heatmode."""
+    client.heatmode = heat_mode
+
+    await common.async_set_hvac_mode(hass, HVAC_SETTINGS[heat_mode], ENTITY_CLIMATE)
+    await client.new_data_cb()
+    await hass.async_block_till_done()
+
+    return hass.states.get(ENTITY_CLIMATE)
+
+
+async def _patch_spa_heatstate(hass, config_entry, heat_state, client):
+    """Patch the heatmode."""
+    client.get_heatstate.return_value = heat_state
+
+    await common.async_set_hvac_mode(hass, HVAC_SETTINGS[heat_state], ENTITY_CLIMATE)
+    await client.new_data_cb()
+    await hass.async_block_till_done()
+
+    return hass.states.get(ENTITY_CLIMATE)

tests/components/balboa/test_config_flow.py

@@ -1,5 +1,5 @@
 """Test the Balboa Spa Client config flow."""
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.balboa.const import CONF_SYNC_TIME, DOMAIN
@@ -12,8 +12,6 @@ from homeassistant.data_entry_flow import (
     RESULT_TYPE_FORM,
 )
 
-from . import BalboaMock
-
 from tests.common import MockConfigEntry
 
 TEST_DATA = {
@@ -22,7 +20,7 @@ TEST_DATA = {
 TEST_ID = "FakeBalboa"
 
 
-async def test_form(hass: HomeAssistant) -> None:
+async def test_form(hass: HomeAssistant, client: MagicMock) -> None:
     """Test we get the form."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -31,23 +29,8 @@ async def test_form(hass: HomeAssistant) -> None:
     assert result["errors"] == {}
 
     with patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.connect",
+        "homeassistant.components.balboa.config_flow.BalboaSpaWifi",
-        new=BalboaMock.connect,
+        return_value=client,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.disconnect",
-        new=BalboaMock.disconnect,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.listen",
-        new=BalboaMock.listen,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.send_mod_ident_req",
-        new=BalboaMock.send_mod_ident_req,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.send_panel_req",
-        new=BalboaMock.send_panel_req,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.spa_configured",
-        new=BalboaMock.spa_configured,
     ), patch(
         "homeassistant.components.balboa.async_setup_entry",
         return_value=True,
@@ -63,19 +46,17 @@ async def test_form(hass: HomeAssistant) -> None:
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-async def test_form_cannot_connect(hass: HomeAssistant) -> None:
+async def test_form_cannot_connect(hass: HomeAssistant, client: MagicMock) -> None:
     """Test we handle cannot connect error."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
     with patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.connect",
+        "homeassistant.components.balboa.config_flow.BalboaSpaWifi",
-        new=BalboaMock.broken_connect,
+        return_value=client,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.disconnect",
-        new=BalboaMock.disconnect,
     ):
+        client.connect.return_value = False
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             TEST_DATA,
@@ -85,16 +66,17 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None:
     assert result2["errors"] == {"base": "cannot_connect"}
 
 
-async def test_unknown_error(hass: HomeAssistant) -> None:
+async def test_unknown_error(hass: HomeAssistant, client: MagicMock) -> None:
     """Test we handle unknown error."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
     with patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.connect",
+        "homeassistant.components.balboa.config_flow.BalboaSpaWifi",
-        side_effect=Exception,
+        return_value=client,
     ):
+        client.connect.side_effect = Exception("Boom")
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             TEST_DATA,
@@ -104,7 +86,7 @@ async def test_unknown_error(hass: HomeAssistant) -> None:
     assert result2["errors"] == {"base": "unknown"}
 
 
-async def test_already_configured(hass: HomeAssistant) -> None:
+async def test_already_configured(hass: HomeAssistant, client: MagicMock) -> None:
     """Test when provided credentials are already configured."""
     MockConfigEntry(domain=DOMAIN, data=TEST_DATA, unique_id=TEST_ID).add_to_hass(hass)
 
@@ -116,11 +98,8 @@ async def test_already_configured(hass: HomeAssistant) -> None:
     assert result["step_id"] == SOURCE_USER
 
     with patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.connect",
+        "homeassistant.components.balboa.config_flow.BalboaSpaWifi",
-        new=BalboaMock.connect,
+        return_value=client,
-    ), patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi.disconnect",
-        new=BalboaMock.disconnect,
     ), patch(
         "homeassistant.components.balboa.async_setup_entry",
         return_value=True,
@@ -135,25 +114,15 @@ async def test_already_configured(hass: HomeAssistant) -> None:
     assert result2["reason"] == "already_configured"
 
 
-async def test_options_flow(hass):
+async def test_options_flow(hass: HomeAssistant, client: MagicMock) -> None:
     """Test specifying non default settings using options flow."""
     config_entry = MockConfigEntry(domain=DOMAIN, data=TEST_DATA, unique_id=TEST_ID)
     config_entry.add_to_hass(hass)
 
-    # Rather than mocking out 15 or so functions, we just need to mock
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
-    # the entire library, otherwise it will get stuck in a listener and
+    await hass.async_block_till_done()
-    # the various loops in pybalboa.
-    with patch(
-        "homeassistant.components.balboa.config_flow.BalboaSpaWifi",
-        new=BalboaMock,
-    ), patch(
-        "homeassistant.components.balboa.BalboaSpaWifi",
-        new=BalboaMock,
-    ):
-        assert await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
 
-        result = await hass.config_entries.options.async_init(config_entry.entry_id)
+    result = await hass.config_entries.options.async_init(config_entry.entry_id)
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "init"
@@ -164,4 +133,4 @@ async def test_options_flow(hass):
     )
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
-    assert config_entry.options == {CONF_SYNC_TIME: True}
+    assert dict(config_entry.options) == {CONF_SYNC_TIME: True}

tests/components/balboa/test_init.py

@@ -1,18 +1,18 @@
 """Tests of the initialization of the balboa integration."""
 
-from unittest.mock import patch
+from unittest.mock import MagicMock
 
 from homeassistant.components.balboa.const import DOMAIN as BALBOA_DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 
-from . import TEST_HOST, BalboaMock, init_integration
+from . import TEST_HOST, init_integration
 
 from tests.common import MockConfigEntry
 
 
-async def test_setup_entry(hass: HomeAssistant):
+async def test_setup_entry(hass: HomeAssistant, client: MagicMock) -> None:
     """Validate that setup entry also configure the client."""
     config_entry = await init_integration(hass)
 
@@ -23,7 +23,7 @@ async def test_setup_entry(hass: HomeAssistant):
     assert config_entry.state == ConfigEntryState.NOT_LOADED
 
 
-async def test_setup_entry_fails(hass):
+async def test_setup_entry_fails(hass: HomeAssistant, client: MagicMock) -> None:
     """Validate that setup entry also configure the client."""
     config_entry = MockConfigEntry(
         domain=BALBOA_DOMAIN,
@@ -33,11 +33,9 @@ async def test_setup_entry_fails(hass):
     )
     config_entry.add_to_hass(hass)
 
-    with patch(
+    client.connect.return_value = False
-        "homeassistant.components.balboa.BalboaSpaWifi.connect",
+
-        new=BalboaMock.broken_connect,
+    await hass.config_entries.async_setup(config_entry.entry_id)
-    ):
+    await hass.async_block_till_done()
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
 
     assert config_entry.state == ConfigEntryState.SETUP_RETRY