Refactor Ecobee Switch entity

homeassistant/components/ecobee/switch.py

@@ -38,15 +38,16 @@ async def async_setup_entry(
     data = hass.data[DOMAIN]
 
     async_add_entities(
-        EcobeeSwitch(data, index, description)
+        EcobeeSwitchAuxHeatOnly(data, index, description)
         for index in range(len(data.ecobee.thermostats))
-        # if data.ecobee.get_thermostat(index)["settings"]["hasHeatPump"]
         for description in SWITCH_TYPES
+        if data.ecobee.get_thermostat(index)["settings"]["hasHeatPump"]
+        and description.key == AUX_HEAT_ONLY_KEY
     )
 
 
-class EcobeeSwitch(SwitchEntity):
-    """Representation of a ecobee switch."""
+class EcobeeSwitchBase(SwitchEntity):
+    """Representation of a base ecobee switch."""
 
     _attr_has_entity_name = True
 
@@ -57,15 +58,25 @@ def __init__(
         self.data = data
         self.thermostat_index = thermostat_index
         self.thermostat = self.data.ecobee.get_thermostat(self.thermostat_index)
-        self._last_hvac_mode_before_aux_heat = HVACMode.HEAT_COOL
         self.entity_description = description
         self._attr_unique_id = f"{self.thermostat['name']}_{description.key}"
-
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, self.thermostat["identifier"])},
             name=self.thermostat["name"],
         )
 
+
+class EcobeeSwitchAuxHeatOnly(EcobeeSwitchBase):
+    """Representation of a aux_heat_only ecobee switch."""
+
+    def __init__(
+        self, data, thermostat_index, description: SwitchEntityDescription
+    ) -> None:
+        """Initialize the switch."""
+        super().__init__(data, thermostat_index, description)
+
+        self._last_hvac_mode_before_aux_heat = HVACMode.HEAT_COOL
+
     def turn_on(self, **kwargs: Any) -> None:
         """Set the hvacMode to auxHeatOnly."""
         _LOGGER.debug("Setting HVAC mode to auxHeatOnly to turn on aux heat")

tests/components/ecobee/test_repairs.py

@@ -51,8 +51,7 @@ async def test_ecobee_repair_flow(
         domain=DOMAIN,
         issue_id="migrate_notify",
     )
-    issue_count = len(issue_registry.issues)
-    assert issue_count >= 1
+    assert len(issue_registry.issues) == 1
 
     url = RepairsFlowIndexView.url
     resp = await http_client.post(
@@ -77,4 +76,4 @@ async def test_ecobee_repair_flow(
         domain=DOMAIN,
         issue_id="migrate_notify",
     )
-    assert len(issue_registry.issues) == (issue_count - 1)
+    assert len(issue_registry.issues) == 0