Rework drop_connect switch, select and coordinator tests and cleanup fixtures (#107119)

* Refactor drop_connect switch and select tests

* Update coordinator tests, cleanup fixtures

tests/components/drop_connect/conftest.py

@@ -1,177 +0,0 @@
-"""Define fixtures available for all tests."""
-import pytest
-
-from homeassistant.components.drop_connect.const import (
-    CONF_COMMAND_TOPIC,
-    CONF_DATA_TOPIC,
-    CONF_DEVICE_DESC,
-    CONF_DEVICE_ID,
-    CONF_DEVICE_NAME,
-    CONF_DEVICE_OWNER_ID,
-    CONF_DEVICE_TYPE,
-    CONF_HUB_ID,
-    DOMAIN,
-)
-from homeassistant.core import HomeAssistant
-
-from tests.common import MockConfigEntry
-
-
-@pytest.fixture
-def config_entry_hub(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_255",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/255/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/255/#",
-            CONF_DEVICE_DESC: "Hub",
-            CONF_DEVICE_ID: 255,
-            CONF_DEVICE_NAME: "Hub DROP-1_C0FFEE",
-            CONF_DEVICE_TYPE: "hub",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_salt(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_8",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/8/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/8/#",
-            CONF_DEVICE_DESC: "Salt Sensor",
-            CONF_DEVICE_ID: 8,
-            CONF_DEVICE_NAME: "Salt Sensor",
-            CONF_DEVICE_TYPE: "salt",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_leak(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_20",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/20/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/20/#",
-            CONF_DEVICE_DESC: "Leak Detector",
-            CONF_DEVICE_ID: 20,
-            CONF_DEVICE_NAME: "Leak Detector",
-            CONF_DEVICE_TYPE: "leak",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_softener(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_0",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/0/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/0/#",
-            CONF_DEVICE_DESC: "Softener",
-            CONF_DEVICE_ID: 0,
-            CONF_DEVICE_NAME: "Softener",
-            CONF_DEVICE_TYPE: "soft",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_filter(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_4",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/4/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/4/#",
-            CONF_DEVICE_DESC: "Filter",
-            CONF_DEVICE_ID: 4,
-            CONF_DEVICE_NAME: "Filter",
-            CONF_DEVICE_TYPE: "filt",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_protection_valve(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_78",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/78/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/78/#",
-            CONF_DEVICE_DESC: "Protection Valve",
-            CONF_DEVICE_ID: 78,
-            CONF_DEVICE_NAME: "Protection Valve",
-            CONF_DEVICE_TYPE: "pv",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_pump_controller(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_83",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/83/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/83/#",
-            CONF_DEVICE_DESC: "Pump Controller",
-            CONF_DEVICE_ID: 83,
-            CONF_DEVICE_NAME: "Pump Controller",
-            CONF_DEVICE_TYPE: "pc",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )
-
-
-@pytest.fixture
-def config_entry_ro_filter(hass: HomeAssistant):
-    """Config entry version 1 fixture."""
-    return MockConfigEntry(
-        domain=DOMAIN,
-        unique_id="DROP-1_C0FFEE_255",
-        data={
-            CONF_COMMAND_TOPIC: "drop_connect/DROP-1_C0FFEE/95/cmd",
-            CONF_DATA_TOPIC: "drop_connect/DROP-1_C0FFEE/95/#",
-            CONF_DEVICE_DESC: "RO Filter",
-            CONF_DEVICE_ID: 95,
-            CONF_DEVICE_NAME: "RO Filter",
-            CONF_DEVICE_TYPE: "ro",
-            CONF_HUB_ID: "DROP-1_C0FFEE",
-            CONF_DEVICE_OWNER_ID: "DROP-1_C0FFEE_255",
-        },
-        version=1,
-    )

tests/components/drop_connect/test_coordinator.py

@@ -1,74 +1,65 @@
 """Test DROP coordinator."""
-from homeassistant.components.drop_connect.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
-from homeassistant.setup import async_setup_component
 
-from .common import TEST_DATA_HUB, TEST_DATA_HUB_RESET, TEST_DATA_HUB_TOPIC
+from .common import (
+    TEST_DATA_HUB,
+    TEST_DATA_HUB_RESET,
+    TEST_DATA_HUB_TOPIC,
+    config_entry_hub,
+)
 
 from tests.common import async_fire_mqtt_message
 from tests.typing import MqttMockHAClient
 
 
-async def test_bad_json(
+async def test_bad_json(hass: HomeAssistant, mqtt_mock: MqttMockHAClient) -> None:
-    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
-) -> None:
     """Test bad JSON."""
-    config_entry_hub.add_to_hass(hass)
+    entry = config_entry_hub()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
 
     current_flow_sensor_name = "sensor.hub_drop_1_c0ffee_water_flow_rate"
-    hass.states.async_set(current_flow_sensor_name, STATE_UNKNOWN)
+    assert hass.states.get(current_flow_sensor_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, "{BAD JSON}")
     await hass.async_block_till_done()
-
+    assert hass.states.get(current_flow_sensor_name).state == STATE_UNKNOWN
-    current_flow_sensor = hass.states.get(current_flow_sensor_name)
-    assert current_flow_sensor
-    assert current_flow_sensor.state == STATE_UNKNOWN
 
 
-async def test_unload(
+async def test_unload(hass: HomeAssistant, mqtt_mock: MqttMockHAClient) -> None:
-    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
-) -> None:
     """Test entity unload."""
     # Load the hub device
-    config_entry_hub.add_to_hass(hass)
+    entry = config_entry_hub()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
 
     current_flow_sensor_name = "sensor.hub_drop_1_c0ffee_water_flow_rate"
-    hass.states.async_set(current_flow_sensor_name, STATE_UNKNOWN)
+    assert hass.states.get(current_flow_sensor_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
+    assert hass.states.get(current_flow_sensor_name).state == "0.0"
+
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
     await hass.async_block_till_done()
 
-    current_flow_sensor = hass.states.get(current_flow_sensor_name)
+    assert hass.states.get(current_flow_sensor_name).state == "5.77"
-    assert current_flow_sensor
-    assert round(float(current_flow_sensor.state), 1) == 5.8
 
     # Unload the device
-    await hass.config_entries.async_unload(config_entry_hub.entry_id)
+    await hass.config_entries.async_unload(entry.entry_id)
-    await hass.async_block_till_done()
+    assert entry.state is ConfigEntryState.NOT_LOADED
-
-    assert config_entry_hub.state is ConfigEntryState.NOT_LOADED
 
     # Verify sensor is unavailable
-    current_flow_sensor = hass.states.get(current_flow_sensor_name)
+    assert hass.states.get(current_flow_sensor_name).state == STATE_UNAVAILABLE
-    assert current_flow_sensor
-    assert current_flow_sensor.state == STATE_UNAVAILABLE
 
 
-async def test_no_mqtt(hass: HomeAssistant, config_entry_hub) -> None:
+async def test_no_mqtt(hass: HomeAssistant) -> None:
     """Test no MQTT."""
-    config_entry_hub.add_to_hass(hass)
+    entry = config_entry_hub()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id) is False
 
     protect_mode_select_name = "select.hub_drop_1_c0ffee_protect_mode"
-    protect_mode_select = hass.states.get(protect_mode_select_name)
+    assert hass.states.get(protect_mode_select_name) is None
-    assert protect_mode_select is None

tests/components/drop_connect/test_select.py

@@ -1,6 +1,5 @@
 """Test DROP select entities."""
 
-from homeassistant.components.drop_connect.const import DOMAIN
 from homeassistant.components.select import (
     ATTR_OPTION,
     ATTR_OPTIONS,
@@ -9,21 +8,23 @@ from homeassistant.components.select import (
 )
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
-from homeassistant.setup import async_setup_component
 
-from .common import TEST_DATA_HUB, TEST_DATA_HUB_RESET, TEST_DATA_HUB_TOPIC
+from .common import (
+    TEST_DATA_HUB,
+    TEST_DATA_HUB_RESET,
+    TEST_DATA_HUB_TOPIC,
+    config_entry_hub,
+)
 
 from tests.common import async_fire_mqtt_message
 from tests.typing import MqttMockHAClient
 
 
-async def test_selects_hub(
+async def test_selects_hub(hass: HomeAssistant, mqtt_mock: MqttMockHAClient) -> None:
-    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
-) -> None:
     """Test DROP binary sensors for hubs."""
-    config_entry_hub.add_to_hass(hass)
+    entry = config_entry_hub()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
 
     protect_mode_select_name = "select.hub_drop_1_c0ffee_protect_mode"
     protect_mode_select = hass.states.get(protect_mode_select_name)
@@ -36,6 +37,14 @@ async def test_selects_hub(
 
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
+    protect_mode_select = hass.states.get(protect_mode_select_name)
+    assert protect_mode_select
+    assert protect_mode_select.attributes.get(ATTR_OPTIONS) == [
+        "away",
+        "home",
+        "schedule",
+    ]
+
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
     await hass.async_block_till_done()
 
@@ -43,6 +52,7 @@ async def test_selects_hub(
     assert protect_mode_select
     assert protect_mode_select.state == "home"
 
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SELECT_DOMAIN,
         SERVICE_SELECT_OPTION,
@@ -50,7 +60,9 @@ async def test_selects_hub(
         blocking=True,
     )
     await hass.async_block_till_done()
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
+    # Simulate response of the device
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
 

tests/components/drop_connect/test_switch.py

@@ -1,6 +1,5 @@
 """Test DROP switch entities."""
 
-from homeassistant.components.drop_connect.const import DOMAIN
 from homeassistant.components.switch import (
     DOMAIN as SWITCH_DOMAIN,
     SERVICE_TURN_OFF,
@@ -8,7 +7,6 @@ from homeassistant.components.switch import (
 )
 from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
-from homeassistant.setup import async_setup_component
 
 from .common import (
     TEST_DATA_FILTER,
@@ -23,253 +21,251 @@ from .common import (
     TEST_DATA_SOFTENER,
     TEST_DATA_SOFTENER_RESET,
     TEST_DATA_SOFTENER_TOPIC,
+    config_entry_filter,
+    config_entry_hub,
+    config_entry_protection_valve,
+    config_entry_softener,
 )
 
 from tests.common import async_fire_mqtt_message
 from tests.typing import MqttMockHAClient
 
 
-async def test_switches_hub(
+async def test_switches_hub(hass: HomeAssistant, mqtt_mock: MqttMockHAClient) -> None:
-    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
-) -> None:
     """Test DROP switches for hubs."""
-    config_entry_hub.add_to_hass(hass)
+    entry = config_entry_hub()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
 
     water_supply_switch_name = "switch.hub_drop_1_c0ffee_water_supply"
-    hass.states.async_set(water_supply_switch_name, STATE_UNKNOWN)
+    assert hass.states.get(water_supply_switch_name).state == STATE_UNKNOWN
     bypass_switch_name = "switch.hub_drop_1_c0ffee_treatment_bypass"
-    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
+    assert hass.states.get(bypass_switch_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
+    assert hass.states.get(water_supply_switch_name).state == STATE_OFF
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
+
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_ON
-    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    assert water_supply_switch
-    assert water_supply_switch.state == STATE_ON
 
     # Test switch turn off method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
         {ATTR_ENTITY_ID: water_supply_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the hub
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_OFF
-    water_supply_switch = hass.states.get(water_supply_switch_name)
-    assert water_supply_switch
-    assert water_supply_switch.state == STATE_OFF
 
     # Test switch turn on method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
         {ATTR_ENTITY_ID: water_supply_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the hub
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_ON
-    water_supply_switch = hass.states.get(water_supply_switch_name)
-    assert water_supply_switch
-    assert water_supply_switch.state == STATE_ON
-
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_OFF
 
     # Test switch turn on method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_ON
 
     # Test switch turn off method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_OFF
 
 
 async def test_switches_protection_valve(
-    hass: HomeAssistant, config_entry_protection_valve, mqtt_mock: MqttMockHAClient
+    hass: HomeAssistant, mqtt_mock: MqttMockHAClient
 ) -> None:
     """Test DROP switches for protection valves."""
-    config_entry_protection_valve.add_to_hass(hass)
+    entry = config_entry_protection_valve()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
+
+    water_supply_switch_name = "switch.protection_valve_water_supply"
+    assert hass.states.get(water_supply_switch_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(
         hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE_RESET
     )
     await hass.async_block_till_done()
+    assert hass.states.get(water_supply_switch_name).state == STATE_OFF
+
     async_fire_mqtt_message(
         hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE
     )
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_ON
-    water_supply_switch_name = "switch.protection_valve_water_supply"
-    hass.states.async_set(water_supply_switch_name, STATE_UNKNOWN)
 
     # Test switch turn off method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
         {ATTR_ENTITY_ID: water_supply_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(
         hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE_RESET
     )
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_OFF
-    water_supply_switch = hass.states.get(water_supply_switch_name)
-    assert water_supply_switch
-    assert water_supply_switch.state == STATE_OFF
 
     # Test switch turn on method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
         {ATTR_ENTITY_ID: water_supply_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(
         hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE
     )
     await hass.async_block_till_done()
-
+    assert hass.states.get(water_supply_switch_name).state == STATE_ON
-    water_supply_switch = hass.states.get(water_supply_switch_name)
-    assert water_supply_switch
-    assert water_supply_switch.state == STATE_ON
 
 
 async def test_switches_softener(
-    hass: HomeAssistant, config_entry_softener, mqtt_mock: MqttMockHAClient
+    hass: HomeAssistant, mqtt_mock: MqttMockHAClient
 ) -> None:
     """Test DROP switches for softeners."""
-    config_entry_softener.add_to_hass(hass)
+    entry = config_entry_softener()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
+
+    bypass_switch_name = "switch.softener_treatment_bypass"
+    assert hass.states.get(bypass_switch_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER_RESET)
     await hass.async_block_till_done()
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
+
     async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    bypass_switch_name = "switch.softener_treatment_bypass"
-    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
 
     # Test switch turn on method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER_RESET)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_ON
 
     # Test switch turn off method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_OFF
 
 
 async def test_switches_filter(
-    hass: HomeAssistant, config_entry_filter, mqtt_mock: MqttMockHAClient
+    hass: HomeAssistant, mqtt_mock: MqttMockHAClient
 ) -> None:
     """Test DROP switches for filters."""
-    config_entry_filter.add_to_hass(hass)
+    entry = config_entry_filter()
-    assert await async_setup_component(hass, DOMAIN, {})
+    entry.add_to_hass(hass)
-    await hass.async_block_till_done()
+    assert await hass.config_entries.async_setup(entry.entry_id)
+
+    bypass_switch_name = "switch.filter_treatment_bypass"
+    assert hass.states.get(bypass_switch_name).state == STATE_UNKNOWN
 
     async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER_RESET)
     await hass.async_block_till_done()
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
+
     async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    bypass_switch_name = "switch.filter_treatment_bypass"
-    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
 
     # Test switch turn on method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_ON,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER_RESET)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_ON
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_ON
 
     # Test switch turn off method.
+    mqtt_mock.async_publish.reset_mock()
     await hass.services.async_call(
         SWITCH_DOMAIN,
         SERVICE_TURN_OFF,
         {ATTR_ENTITY_ID: bypass_switch_name},
         blocking=True,
     )
+    assert len(mqtt_mock.async_publish.mock_calls) == 1
 
     # Simulate response from the device
     async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER)
     await hass.async_block_till_done()
-
+    assert hass.states.get(bypass_switch_name).state == STATE_OFF
-    bypass_switch = hass.states.get(bypass_switch_name)
-    assert bypass_switch
-    assert bypass_switch.state == STATE_OFF