From 9ee0d8fefe2173b7aa5d9e36704f559362681a76 Mon Sep 17 00:00:00 2001
From: Maciej Bieniek <bieniu@users.noreply.github.com>
Date: Mon, 6 Sep 2021 15:30:03 +0200
Subject: [PATCH] Fix xiaomi miio Air Quality Monitor initialization (#55773)

---
 homeassistant/components/xiaomi_miio/const.py |  7 +++
 .../components/xiaomi_miio/sensor.py          | 59 ++++++++++---------
 2 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/homeassistant/components/xiaomi_miio/const.py b/homeassistant/components/xiaomi_miio/const.py
index b63143c0f41..29c740cf800 100644
--- a/homeassistant/components/xiaomi_miio/const.py
+++ b/homeassistant/components/xiaomi_miio/const.py
@@ -118,6 +118,13 @@ MODEL_AIRQUALITYMONITOR_B1 = "cgllc.airmonitor.b1"
 MODEL_AIRQUALITYMONITOR_S1 = "cgllc.airmonitor.s1"
 MODEL_AIRQUALITYMONITOR_CGDN1 = "cgllc.airm.cgdn1"
 
+MODELS_AIR_QUALITY_MONITOR = [
+    MODEL_AIRQUALITYMONITOR_V1,
+    MODEL_AIRQUALITYMONITOR_B1,
+    MODEL_AIRQUALITYMONITOR_S1,
+    MODEL_AIRQUALITYMONITOR_CGDN1,
+]
+
 # Light Models
 MODELS_LIGHT_EYECARE = ["philips.light.sread1"]
 MODELS_LIGHT_CEILING = ["philips.light.ceiling", "philips.light.zyceiling"]
diff --git a/homeassistant/components/xiaomi_miio/sensor.py b/homeassistant/components/xiaomi_miio/sensor.py
index 96bcdf9145d..7e2a5d230cd 100644
--- a/homeassistant/components/xiaomi_miio/sensor.py
+++ b/homeassistant/components/xiaomi_miio/sensor.py
@@ -68,6 +68,7 @@ from .const import (
     MODEL_FAN_ZA1,
     MODEL_FAN_ZA3,
     MODEL_FAN_ZA4,
+    MODELS_AIR_QUALITY_MONITOR,
     MODELS_HUMIDIFIER_MIIO,
     MODELS_HUMIDIFIER_MIOT,
     MODELS_HUMIDIFIER_MJJSQ,
@@ -380,23 +381,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         host = config_entry.data[CONF_HOST]
         token = config_entry.data[CONF_TOKEN]
         model = config_entry.data[CONF_MODEL]
-        device = hass.data[DOMAIN][config_entry.entry_id].get(KEY_DEVICE)
-        sensors = []
+
         if model in (MODEL_FAN_ZA1, MODEL_FAN_ZA3, MODEL_FAN_ZA4, MODEL_FAN_P5):
             return
-        if model in MODEL_TO_SENSORS_MAP:
-            sensors = MODEL_TO_SENSORS_MAP[model]
-        elif model in MODELS_HUMIDIFIER_MIOT:
-            sensors = HUMIDIFIER_MIOT_SENSORS
-        elif model in MODELS_HUMIDIFIER_MJJSQ:
-            sensors = HUMIDIFIER_MJJSQ_SENSORS
-        elif model in MODELS_HUMIDIFIER_MIIO:
-            sensors = HUMIDIFIER_MIIO_SENSORS
-        elif model in MODELS_PURIFIER_MIIO:
-            sensors = PURIFIER_MIIO_SENSORS
-        elif model in MODELS_PURIFIER_MIOT:
-            sensors = PURIFIER_MIOT_SENSORS
-        else:
+
+        if model in MODELS_AIR_QUALITY_MONITOR:
             unique_id = config_entry.unique_id
             name = config_entry.title
             _LOGGER.debug("Initializing with host %s (token %s...)", host, token[:5])
@@ -408,19 +397,35 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
                     name, device, config_entry, unique_id, description
                 )
             )
-        for sensor, description in SENSOR_TYPES.items():
-            if sensor not in sensors:
-                continue
-            entities.append(
-                XiaomiGenericSensor(
-                    f"{config_entry.title} {description.name}",
-                    device,
-                    config_entry,
-                    f"{sensor}_{config_entry.unique_id}",
-                    hass.data[DOMAIN][config_entry.entry_id][KEY_COORDINATOR],
-                    description,
+        else:
+            device = hass.data[DOMAIN][config_entry.entry_id][KEY_DEVICE]
+            sensors = []
+            if model in MODEL_TO_SENSORS_MAP:
+                sensors = MODEL_TO_SENSORS_MAP[model]
+            elif model in MODELS_HUMIDIFIER_MIOT:
+                sensors = HUMIDIFIER_MIOT_SENSORS
+            elif model in MODELS_HUMIDIFIER_MJJSQ:
+                sensors = HUMIDIFIER_MJJSQ_SENSORS
+            elif model in MODELS_HUMIDIFIER_MIIO:
+                sensors = HUMIDIFIER_MIIO_SENSORS
+            elif model in MODELS_PURIFIER_MIIO:
+                sensors = PURIFIER_MIIO_SENSORS
+            elif model in MODELS_PURIFIER_MIOT:
+                sensors = PURIFIER_MIOT_SENSORS
+
+            for sensor, description in SENSOR_TYPES.items():
+                if sensor not in sensors:
+                    continue
+                entities.append(
+                    XiaomiGenericSensor(
+                        f"{config_entry.title} {description.name}",
+                        device,
+                        config_entry,
+                        f"{sensor}_{config_entry.unique_id}",
+                        hass.data[DOMAIN][config_entry.entry_id][KEY_COORDINATOR],
+                        description,
+                    )
                 )
-            )
 
     async_add_entities(entities)