Fix #161

Limych · Jul 3, 2024 · 6e409a2 · 6e409a2
1 parent eb8b744
commit 6e409a2
Show file tree

Hide file tree

Showing 2 changed files with 120 additions and 34 deletions.
diff --git a/custom_components/apparent_temperature/sensor.py b/custom_components/apparent_temperature/sensor.py
@@ -151,6 +151,51 @@ def extra_state_attributes(self):
             ATTR_WIND_SPEED_SOURCE_VALUE: self._wind_val,
         }
 
+    def _setup_sources(self) -> list[str]:
+        """Set sources for entity and return list of sources to track."""
+        entities = set()
+        for entity_id in self._sources:
+            state: State = self.hass.states.get(entity_id)
+            domain = split_entity_id(state.entity_id)[0]
+            device_class = state.attributes.get(ATTR_DEVICE_CLASS)
+            unit_of_measurement = state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
+
+            if domain == WEATHER_DOMAIN:
+                self._temp = entity_id
+                self._humd = entity_id
+                self._wind = entity_id
+                entities.add(entity_id)
+            elif domain == CLIMATE_DOMAIN:
+                self._temp = entity_id
+                self._humd = entity_id
+                entities.add(entity_id)
+            elif (
+                device_class == SensorDeviceClass.TEMPERATURE
+                or unit_of_measurement in UnitOfTemperature
+            ):
+                self._temp = entity_id
+                entities.add(entity_id)
+            elif (
+                device_class == SensorDeviceClass.HUMIDITY
+                or unit_of_measurement == PERCENTAGE
+            ):
+                self._humd = entity_id
+                entities.add(entity_id)
+            elif unit_of_measurement in UnitOfSpeed:
+                self._wind = entity_id
+                entities.add(entity_id)
+            elif entity_id.find("temperature") >= 0:
+                self._temp = entity_id
+                entities.add(entity_id)
+            elif entity_id.find("humidity") >= 0:
+                self._humd = entity_id
+                entities.add(entity_id)
+            elif entity_id.find("wind") >= 0:
+                self._wind = entity_id
+                entities.add(entity_id)
+
+        return list(entities)
+
     async def async_added_to_hass(self):
         """Register callbacks."""
 
@@ -164,41 +209,8 @@ def sensor_state_listener(event: Event[EventStateChangedData]) -> None:
         @callback
         def sensor_startup(event):
             """Update entity on startup."""
-            entities = set()
-            for entity_id in self._sources:
-                state: State = self.hass.states.get(entity_id)
-                domain = split_entity_id(state.entity_id)[0]
-                device_class = state.attributes.get(ATTR_DEVICE_CLASS)
-                unit_of_measurement = state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
-
-                if (
-                    device_class == SensorDeviceClass.TEMPERATURE
-                    or domain in (WEATHER_DOMAIN, CLIMATE_DOMAIN)
-                    or unit_of_measurement in UnitOfTemperature
-                    or entity_id.find("temperature") >= 0
-                ):
-                    self._temp = entity_id
-                    entities.add(entity_id)
-
-                if (
-                    device_class == SensorDeviceClass.HUMIDITY
-                    or domain in (WEATHER_DOMAIN, CLIMATE_DOMAIN)
-                    or unit_of_measurement == PERCENTAGE
-                    or entity_id.find("humidity") >= 0
-                ):
-                    self._humd = entity_id
-                    entities.add(entity_id)
-
-                if (
-                    domain == WEATHER_DOMAIN
-                    or unit_of_measurement in UnitOfSpeed
-                    or entity_id.find("wind") >= 0
-                ):
-                    self._wind = entity_id
-                    entities.add(entity_id)
-
             async_track_state_change_event(
-                self.hass, list(entities), sensor_state_listener
+                self.hass, self._setup_sources(), sensor_state_listener
             )
 
             self.async_schedule_update_ha_state(True)  # Force first update

diff --git a/tests/test_sensor.py b/tests/test_sensor.py
@@ -124,13 +124,87 @@ async def async_setup_test_entities(hass: HomeAssistant):
                                 "value_template": "{{ 10 }}",
                                 "availability_template": "{{ False }}",
                             },
+                            "test_temperature_no_unit": {
+                                "value_template": "{{ 20 }}",
+                            },
+                            "test_humidity_no_unit": {
+                                "value_template": "{{ 40 }}",
+                            },
+                            "test_wind_speed_no_unit": {
+                                "value_template": "{{ 10 }}",
+                            },
                         },
                     },
                 ],
             },
         )
     await hass.async_block_till_done()
 
+    assert await async_setup_component(
+        hass,
+        "climate",
+        {
+            "climate": {
+                "platform": "generic_thermostat",
+                "name": "test_climate",
+                "heater": "switch.study_heater",
+                "target_sensor": "sensor.test_temperature",
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+
+async def test__setup_sources(hass: HomeAssistant):
+    """Test setup sensor sources."""
+    await async_setup_test_entities(hass)
+    entity = ApparentTemperatureSensor(None, TEST_NAME, TEST_SOURCES)
+    entity.hass = hass
+
+    assert entity._temp is None
+    assert entity._humd is None
+    assert entity._wind is None
+
+    entity._temp = entity._humd = entity._wind = None
+    entity._sources = ["weather.test_monitored"]
+    entity._setup_sources()
+
+    assert entity._temp == "weather.test_monitored"
+    assert entity._humd == "weather.test_monitored"
+    assert entity._wind == "weather.test_monitored"
+
+    entity._temp = entity._humd = entity._wind = None
+    entity._sources = ["climate.test_climate"]
+    entity._setup_sources()
+
+    assert entity._temp == "climate.test_climate"
+    assert entity._humd == "climate.test_climate"
+    assert entity._wind is None
+
+    entity._temp = entity._humd = entity._wind = None
+    entity._sources = [
+        "sensor.test_temperature",
+        "sensor.test_humidity",
+        "sensor.test_wind_speed",
+    ]
+    entity._setup_sources()
+
+    assert entity._temp == "sensor.test_temperature"
+    assert entity._humd == "sensor.test_humidity"
+    assert entity._wind == "sensor.test_wind_speed"
+
+    entity._temp = entity._humd = entity._wind = None
+    entity._sources = [
+        "sensor.test_temperature_no_unit",
+        "sensor.test_humidity_no_unit",
+        "sensor.test_wind_speed_no_unit",
+    ]
+    entity._setup_sources()
+
+    assert entity._temp == "sensor.test_temperature_no_unit"
+    assert entity._humd == "sensor.test_humidity_no_unit"
+    assert entity._wind == "sensor.test_wind_speed_no_unit"
+
 
 async def test_entity_initialization(hass: HomeAssistant):
     """Test sensor initialization."""