Persist extra battery config on initial setup and migrate old entries (#255)

Extra batteries were detected on the first setup, but they weren't written into the config due to oversight - only reconfigure menu contains schema for extra batteries. This means that initial configuration did not have extra battery keys and was relying on runtime detection which could fail if the message does not arrive before this check. User had to explicitly reconfigure it for it to persist. This PR fixes this and adds migration for old entries to contain extra batteries key.

Also it removes second extra battery slot from River 3 Plus as it only supports a single battery.

Author: GnoX

custom_components/ef_ble/__init__.py

@@ -22,6 +22,7 @@
     CONF_BLUEZ_START_NOTIFY,
     CONF_COLLECT_PACKETS_AMOUNT,
     CONF_CONNECTION_TIMEOUT,
+    CONF_EXTRA_BATTERY,
     CONF_PACKET_VERSION,
     CONF_UPDATE_PERIOD,
     CONF_USER_ID,
@@ -188,8 +189,8 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         config_entry.minor_version,
     )
 
-    match config_entry.version, config_entry.minor_version:
-        case 1, 0:
+    if config_entry.version < 2:
+        if config_entry.minor_version < 1:
             address = config_entry.data.get(CONF_ADDRESS)
             device_reg = dr.async_get(hass)
             device_entry = device_reg.async_get_device(identifiers={(DOMAIN, address)})
@@ -212,6 +213,13 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
 
             hass.config_entries.async_update_entry(config_entry, minor_version=1)
 
+        if config_entry.minor_version < 2:
+            data = {**config_entry.data}
+            data.setdefault(CONF_EXTRA_BATTERY, [])
+            hass.config_entries.async_update_entry(
+                config_entry, data=data, minor_version=2
+            )
+
     return True
 
 

custom_components/ef_ble/config_flow.py

@@ -6,7 +6,7 @@
 import base64
 import enum
 import logging
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
 from functools import cached_property
 from typing import Any, ClassVar, cast
 
@@ -88,7 +88,7 @@ class EFBLEConfigFlow(ConfigFlow, domain=DOMAIN):
     """EcoFlow BLE ConfigFlow"""
 
     VERSION = 1
-    MINOR_VERSION = 1
+    MINOR_VERSION = 2
 
     CONNECTION_CLASS = CONN_CLASS_LOCAL_PUSH
 
@@ -385,9 +385,14 @@ def _create_entry(self, user_input: dict[str, Any], device: eflib.DeviceBase):
         entry_data["local_name"] = self._local_names.get(device.address, None)
         entry_data.pop("login", None)
 
+        if CONF_EXTRA_BATTERY not in entry_data:
+            entry_data[CONF_EXTRA_BATTERY] = _find_enabled_batteries(
+                device, range(1, 6)
+            )
+
         return self.async_create_entry(title=device.name, data=entry_data)
 
-    def _check_user_id(self, user_id: str):
+    def _check_user_id(self, user_id: str) -> dict[str, str] | None:
         try:
             int(user_id.strip())
         except ValueError:
@@ -758,11 +763,7 @@ def extra_battery(
         extra_batteries_default = (
             extra_battery_conf
             if extra_battery_conf is not None
-            else [
-                str(i)
-                for i in available_battery_slots
-                if getattr(device, f"battery_{i}_enabled", False)
-            ]
+            else _find_enabled_batteries(device, available_battery_slots)
         )
 
         extra_battery_labels = {
@@ -789,3 +790,7 @@ def extra_battery(
 
 def schema_builder():
     return _SchemaBuilder()
+
+
+def _find_enabled_batteries(device: eflib.DeviceBase, slots: Iterable[int]):
+    return [str(i) for i in slots if getattr(device, f"battery_{i}_enabled", False)]

custom_components/ef_ble/eflib/devices/river3_plus.py

@@ -27,11 +27,6 @@ class Device(river3.Device):
     battery_1_cell_temperature = pb_field(pb.plug_in_info_dcp_resv, resv_temperature)
     battery_1_sn = pb_field(pb.plug_in_info_dcp_sn)
 
-    battery_2_enabled = pb_field(pb.plug_in_info_dcp2_in_flag)
-    battery_2_battery_level = pb_field(pb.plug_in_info_dcp2_resv, resv_soc)
-    battery_2_cell_temperature = pb_field(pb.plug_in_info_dcp2_resv, resv_temperature)
-    battery_2_sn = pb_field(pb.plug_in_info_dcp2_sn)
-
     led_mode = pb_field(river3.pb.led_mode, LedMode.from_value)
 
     @controls.select(led_mode, options=LedMode)