Implement automatic migration for Growatt Server DEFAULT_PLANT_ID entries

homeassistant/components/growatt_server/__init__.py

@@ -41,15 +41,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-def get_device_list_classic(
-    api: growattServer.GrowattApi, config: Mapping[str, str]
-) -> tuple[list[dict[str, str]], str]:
-    """Retrieve the device list for the selected plant."""
-    plant_id = config[CONF_PLANT_ID]
-
-    # Log in to api and fetch first plant if no plant id is defined.
+def _login_classic_api(
+    api: growattServer.GrowattApi, username: str, password: str
+) -> dict:
+    """Log in to Classic API and return user info."""
     try:
-        login_response = api.login(config[CONF_USERNAME], config[CONF_PASSWORD])
+        login_response = api.login(username, password)
     except (RequestException, JSONDecodeError) as ex:
         raise ConfigEntryError(
             f"Error communicating with Growatt API during login: {ex}"
@@ -62,21 +59,45 @@ def get_device_list_classic(
             raise ConfigEntryAuthFailed("Username, Password or URL may be incorrect!")
         raise ConfigEntryError(f"Growatt login failed: {msg}")
 
-    user_id = login_response["user"]["id"]
+    return login_response
 
-    # Legacy support: DEFAULT_PLANT_ID ("0") triggers auto-selection of first plant.
-    # Modern config flow always sets a specific plant_id, but old config entries
-    # from earlier versions may still have plant_id="0".
-    if plant_id == DEFAULT_PLANT_ID:
-        try:
-            plant_info = api.plant_list(user_id)
-        except (RequestException, JSONDecodeError) as ex:
-            raise ConfigEntryError(
-                f"Error communicating with Growatt API during plant list: {ex}"
-            ) from ex
-        if not plant_info or "data" not in plant_info or not plant_info["data"]:
-            raise ConfigEntryError("No plants found for this account.")
-        plant_id = plant_info["data"][0]["plantId"]
+
+def _get_first_plant_id(
+    api: growattServer.GrowattApi, username: str, password: str
+) -> str:
+    """Get the first available plant_id for the authenticated user in Classic API."""
+    login_response = _login_classic_api(api, username, password)
+
+    # Safely extract user_id with validation
+    user_data = login_response.get("user")
+    if not user_data or "id" not in user_data:
+        raise ConfigEntryError("Login response missing required user information.")
+
+    user_id = user_data["id"]
+
+    try:
+        plant_info = api.plant_list(user_id)
+    except (RequestException, JSONDecodeError) as ex:
+        raise ConfigEntryError(
+            f"Error communicating with Growatt API during plant list: {ex}"
+        ) from ex
+
+    if not plant_info or "data" not in plant_info or not plant_info["data"]:
+        raise ConfigEntryError("No plants found for this account.")
+
+    return plant_info["data"][0]["plantId"]
+
+
+def get_device_list_classic(
+    api: growattServer.GrowattApi, config: Mapping[str, str]
+) -> tuple[list[dict[str, str]], str]:
+    """Retrieve the device list for the selected plant."""
+    plant_id = config[CONF_PLANT_ID]
+    username = config[CONF_USERNAME]
+    password = config[CONF_PASSWORD]
+
+    # Login to Classic API
+    _login_classic_api(api, username, password)
 
     # Get a list of devices for specified plant to add sensors for.
     try:
@@ -94,9 +115,9 @@ def get_device_list_v1(
 ) -> tuple[list[dict[str, str]], str]:
     """Device list logic for Open API V1.
 
-    Note: Plant selection (including auto-selection if only one plant exists)
-    is handled in the config flow before this function is called. This function
-    only fetches devices for the already-selected plant_id.
+    Plant selection is handled in the config flow before this function is called.
+    This function expects a specific plant_id and fetches devices for that plant.
+
     """
     plant_id = config[CONF_PLANT_ID]
     try:
@@ -184,6 +205,38 @@ async def async_setup_entry(
     else:
         raise ConfigEntryError("Unknown authentication type in config entry.")
 
+    # Handle migration for entries with DEFAULT_PLANT_ID
+    if config.get(CONF_PLANT_ID) == DEFAULT_PLANT_ID:
+        # Resolve DEFAULT_PLANT_ID to actual plant_id for Classic API
+        if api_version == "classic":
+            try:
+                first_plant_id = await hass.async_add_executor_job(
+                    _get_first_plant_id,
+                    api,
+                    config[CONF_USERNAME],
+                    config[CONF_PASSWORD],
+                )
+            except (RequestException, JSONDecodeError) as ex:
+                raise ConfigEntryError(
+                    f"Error resolving plant_id during migration: {ex}"
+                ) from ex
+
+            # Update config with real plant_id
+            new_data = dict(config_entry.data)
+            new_data[CONF_PLANT_ID] = first_plant_id
+            hass.config_entries.async_update_entry(config_entry, data=new_data)
+            config = config_entry.data  # Refresh config with updated data
+
+            _LOGGER.info(
+                "Migrated config entry to use specific plant_id '%s' instead of default plant selection",
+                first_plant_id,
+            )
+        else:
+            # V1 API should not have DEFAULT_PLANT_ID, but log and continue
+            _LOGGER.warning(
+                "V1 API config entry unexpectedly has DEFAULT_PLANT_ID, this may cause API errors"
+            )
+
     devices, plant_id = await hass.async_add_executor_job(
         get_device_list, api, config, api_version
     )

tests/components/growatt_server/conftest.py

@@ -265,9 +265,10 @@ def mock_config_entry_classic() -> MockConfigEntry:
 def mock_config_entry_classic_default_plant() -> MockConfigEntry:
     """Return a mocked config entry for Classic API with DEFAULT_PLANT_ID.
 
-    This config entry uses plant_id="0" which triggers auto-plant-selection logic
-    in the Classic API path. This is legacy support for old config entries that
-    didn't have a specific plant_id set during initial configuration.
+    This config entry uses plant_id="0" which triggers migration logic in
+    async_setup_entry to resolve to the actual plant_id. This is legacy support
+    for old config entries that didn't have a specific plant_id set during initial
+    configuration.
     """
     return MockConfigEntry(
         domain=DOMAIN,
@@ -276,7 +277,7 @@ def mock_config_entry_classic_default_plant() -> MockConfigEntry:
             CONF_USERNAME: "test_user",
             CONF_PASSWORD: "test_password",
             CONF_URL: "https://server.growatt.com/",
-            CONF_PLANT_ID: DEFAULT_PLANT_ID,  # "0" triggers auto-selection
+            CONF_PLANT_ID: DEFAULT_PLANT_ID,  # "0" - should trigger migration
         },
         unique_id="plant_default",
     )

tests/components/growatt_server/test_init.py

@@ -13,6 +13,8 @@
     AUTH_API_TOKEN,
     AUTH_PASSWORD,
     CONF_AUTH_TYPE,
+    CONF_PLANT_ID,
+    DEFAULT_PLANT_ID,
     DOMAIN,
 )
 from homeassistant.config_entries import ConfigEntryState
@@ -393,3 +395,42 @@ async def test_v1_api_unsupported_device_type(
     assert mock_config_entry.state is ConfigEntryState.LOADED
     # Verify warning was logged for unsupported device
     assert "Device TLX789012 with type 5 not supported in Open API V1" in caplog.text
+
+
+async def test_migrate_default_plant_id(
+    hass: HomeAssistant,
+    mock_growatt_classic_api,
+    mock_config_entry_classic_default_plant: MockConfigEntry,
+) -> None:
+    """Test migration of config entry with DEFAULT_PLANT_ID to actual plant_id."""
+    # Initially has DEFAULT_PLANT_ID
+    assert (
+        mock_config_entry_classic_default_plant.data[CONF_PLANT_ID] == DEFAULT_PLANT_ID
+    )
+
+    # Mock successful login and plant list with specific plant ID
+    mock_growatt_classic_api.login.return_value = {
+        "success": True,
+        "user": {"id": 123456},
+    }
+    mock_growatt_classic_api.plant_list.return_value = {
+        "data": [{"plantId": "MIGRATED_PLANT_123", "plantName": "My Plant"}]
+    }
+    mock_growatt_classic_api.device_list.return_value = [
+        {"deviceSn": "TLX123456", "deviceType": "tlx"}
+    ]
+
+    await setup_integration(hass, mock_config_entry_classic_default_plant)
+
+    # Verify integration loaded successfully
+    assert mock_config_entry_classic_default_plant.state is ConfigEntryState.LOADED
+
+    # Verify config entry was migrated to use the actual plant_id
+    assert (
+        mock_config_entry_classic_default_plant.data[CONF_PLANT_ID]
+        == "MIGRATED_PLANT_123"
+    )
+
+    # Verify other config data remains unchanged
+    assert mock_config_entry_classic_default_plant.data[CONF_AUTH_TYPE] == AUTH_PASSWORD
+    assert mock_config_entry_classic_default_plant.data[CONF_USERNAME] == "test_user"