Fix for justnimbus integration (#99212)

* Fix for justnimbus integration

* Fixed tests and moved const

* fix: added reauth flow

* fix: fixed reauth config flow

* chore: added config_flow reauth test

* chore: Processed PR feedback

---------

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

homeassistant/components/justnimbus/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
 
 from .const import DOMAIN, PLATFORMS
 from .coordinator import JustNimbusCoordinator
@@ -10,7 +11,10 @@ from .coordinator import JustNimbusCoordinator
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up JustNimbus from a config entry."""
+    if "zip_code" in entry.data:
         coordinator = JustNimbusCoordinator(hass=hass, entry=entry)
+    else:
+        raise ConfigEntryAuthFailed()
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator

homeassistant/components/justnimbus/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow for JustNimbus integration."""
 from __future__ import annotations
 
+from collections.abc import Mapping
 import logging
 from typing import Any
 
@@ -12,13 +13,14 @@ from homeassistant.const import CONF_CLIENT_ID
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import config_validation as cv
 
-from .const import DOMAIN
+from .const import CONF_ZIP_CODE, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
 STEP_USER_DATA_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_CLIENT_ID): cv.string,
+        vol.Required(CONF_ZIP_CODE): cv.string,
     },
 )
 
@@ -27,6 +29,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for JustNimbus."""
 
     VERSION = 1
+    reauth_entry: config_entries.ConfigEntry | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -39,10 +42,14 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         errors = {}
 
-        await self.async_set_unique_id(user_input[CONF_CLIENT_ID])
+        unique_id = f"{user_input[CONF_CLIENT_ID]}{user_input[CONF_ZIP_CODE]}"
+        await self.async_set_unique_id(unique_id=unique_id)
+        if not self.reauth_entry:
             self._abort_if_unique_id_configured()
 
-        client = justnimbus.JustNimbusClient(client_id=user_input[CONF_CLIENT_ID])
+        client = justnimbus.JustNimbusClient(
+            client_id=user_input[CONF_CLIENT_ID], zip_code=user_input[CONF_ZIP_CODE]
+        )
         try:
             await self.hass.async_add_executor_job(client.get_data)
         except justnimbus.InvalidClientID:
@@ -53,8 +60,25 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             _LOGGER.exception("Unexpected exception")
             errors["base"] = "unknown"
         else:
+            if not self.reauth_entry:
                 return self.async_create_entry(title="JustNimbus", data=user_input)
+            self.hass.config_entries.async_update_entry(
+                self.reauth_entry, data=user_input, unique_id=unique_id
+            )
+
+            # Reload the config entry otherwise devices will remain unavailable
+            self.hass.async_create_task(
+                self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
+            )
+            return self.async_abort(reason="reauth_successful")
 
         return self.async_show_form(
             step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
         )
+
+    async def async_step_reauth(self, user_input: Mapping[str, Any]) -> FlowResult:
+        """Perform reauth upon an API authentication error."""
+        self.reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+        return await self.async_step_user()

homeassistant/components/justnimbus/const.py

@@ -1,13 +1,14 @@
 """Constants for the JustNimbus integration."""
 
+
 from typing import Final
 
 from homeassistant.const import Platform
 
 DOMAIN = "justnimbus"
 
-VOLUME_FLOW_RATE_LITERS_PER_MINUTE: Final = "L/min"
-
 PLATFORMS = [
     Platform.SENSOR,
 ]
+
+CONF_ZIP_CODE: Final = "zip_code"

homeassistant/components/justnimbus/coordinator.py

@@ -11,7 +11,7 @@ from homeassistant.const import CONF_CLIENT_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DOMAIN
+from .const import CONF_ZIP_CODE, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -27,7 +27,9 @@ class JustNimbusCoordinator(DataUpdateCoordinator[justnimbus.JustNimbusModel]):
             name=DOMAIN,
             update_interval=timedelta(minutes=1),
         )
-        self._client = justnimbus.JustNimbusClient(client_id=entry.data[CONF_CLIENT_ID])
+        self._client = justnimbus.JustNimbusClient(
+            client_id=entry.data[CONF_CLIENT_ID], zip_code=entry.data[CONF_ZIP_CODE]
+        )
 
     async def _async_update_data(self) -> justnimbus.JustNimbusModel:
         """Fetch the latest data from the source."""

homeassistant/components/justnimbus/manifest.json

@@ -5,5 +5,5 @@
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/justnimbus",
   "iot_class": "cloud_polling",
-  "requirements": ["justnimbus==0.6.0"]
+  "requirements": ["justnimbus==0.7.3"]
 }

homeassistant/components/justnimbus/sensor.py

@@ -17,7 +17,6 @@ from homeassistant.const import (
     EntityCategory,
     UnitOfPressure,
     UnitOfTemperature,
-    UnitOfTime,
     UnitOfVolume,
 )
 from homeassistant.core import HomeAssistant
@@ -25,7 +24,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
 from . import JustNimbusCoordinator
-from .const import DOMAIN, VOLUME_FLOW_RATE_LITERS_PER_MINUTE
+from .const import DOMAIN
 from .entity import JustNimbusEntity
 
 
@@ -44,54 +43,20 @@ class JustNimbusEntityDescription(
 
 
 SENSOR_TYPES = (
-    JustNimbusEntityDescription(
-        key="pump_flow",
-        translation_key="pump_flow",
-        icon="mdi:pump",
-        native_unit_of_measurement=VOLUME_FLOW_RATE_LITERS_PER_MINUTE,
-        state_class=SensorStateClass.MEASUREMENT,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.pump_flow,
-    ),
-    JustNimbusEntityDescription(
-        key="drink_flow",
-        translation_key="drink_flow",
-        icon="mdi:water-pump",
-        native_unit_of_measurement=VOLUME_FLOW_RATE_LITERS_PER_MINUTE,
-        state_class=SensorStateClass.MEASUREMENT,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.drink_flow,
-    ),
     JustNimbusEntityDescription(
         key="pump_pressure",
         translation_key="pump_pressure",
+        icon="mdi:water-pump",
         native_unit_of_measurement=UnitOfPressure.BAR,
         device_class=SensorDeviceClass.PRESSURE,
         state_class=SensorStateClass.MEASUREMENT,
         entity_category=EntityCategory.DIAGNOSTIC,
         value_fn=lambda coordinator: coordinator.data.pump_pressure,
     ),
-    JustNimbusEntityDescription(
-        key="pump_starts",
-        translation_key="pump_starts",
-        icon="mdi:restart",
-        state_class=SensorStateClass.MEASUREMENT,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.pump_starts,
-    ),
-    JustNimbusEntityDescription(
-        key="pump_hours",
-        translation_key="pump_hours",
-        icon="mdi:clock",
-        device_class=SensorDeviceClass.DURATION,
-        native_unit_of_measurement=UnitOfTime.HOURS,
-        state_class=SensorStateClass.MEASUREMENT,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.pump_hours,
-    ),
     JustNimbusEntityDescription(
         key="reservoir_temp",
         translation_key="reservoir_temperature",
+        icon="mdi:coolant-temperature",
         native_unit_of_measurement=UnitOfTemperature.CELSIUS,
         device_class=SensorDeviceClass.TEMPERATURE,
         state_class=SensorStateClass.MEASUREMENT,
@@ -104,57 +69,46 @@ SENSOR_TYPES = (
         icon="mdi:car-coolant-level",
         native_unit_of_measurement=UnitOfVolume.LITERS,
         device_class=SensorDeviceClass.VOLUME,
-        state_class=SensorStateClass.MEASUREMENT,
+        state_class=SensorStateClass.TOTAL,
         entity_category=EntityCategory.DIAGNOSTIC,
         value_fn=lambda coordinator: coordinator.data.reservoir_content,
     ),
     JustNimbusEntityDescription(
-        key="total_saved",
+        key="water_saved",
-        translation_key="total_saved",
+        translation_key="water_saved",
         icon="mdi:water-opacity",
         native_unit_of_measurement=UnitOfVolume.LITERS,
         device_class=SensorDeviceClass.VOLUME,
         state_class=SensorStateClass.TOTAL_INCREASING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.total_saved,
+        value_fn=lambda coordinator: coordinator.data.water_saved,
     ),
     JustNimbusEntityDescription(
-        key="total_replenished",
+        key="water_used",
-        translation_key="total_replenished",
+        translation_key="water_used",
-        icon="mdi:water",
-        native_unit_of_measurement=UnitOfVolume.LITERS,
-        device_class=SensorDeviceClass.VOLUME,
-        state_class=SensorStateClass.TOTAL_INCREASING,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.total_replenished,
-    ),
-    JustNimbusEntityDescription(
-        key="error_code",
-        translation_key="error_code",
-        icon="mdi:bug",
-        entity_registry_enabled_default=False,
-        entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.error_code,
-    ),
-    JustNimbusEntityDescription(
-        key="totver",
-        translation_key="total_use",
         icon="mdi:chart-donut",
         native_unit_of_measurement=UnitOfVolume.LITERS,
         device_class=SensorDeviceClass.VOLUME,
         state_class=SensorStateClass.TOTAL_INCREASING,
         entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.totver,
+        value_fn=lambda coordinator: coordinator.data.water_used,
     ),
     JustNimbusEntityDescription(
-        key="reservoir_content_max",
+        key="reservoir_capacity",
-        translation_key="reservoir_content_max",
+        translation_key="reservoir_capacity",
         icon="mdi:waves",
         native_unit_of_measurement=UnitOfVolume.LITERS,
         device_class=SensorDeviceClass.VOLUME,
         state_class=SensorStateClass.TOTAL,
         entity_category=EntityCategory.DIAGNOSTIC,
-        value_fn=lambda coordinator: coordinator.data.reservoir_content_max,
+        value_fn=lambda coordinator: coordinator.data.reservoir_capacity,
+    ),
+    JustNimbusEntityDescription(
+        key="pump_type",
+        translation_key="pump_type",
+        icon="mdi:pump",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda coordinator: coordinator.data.pump_type,
     ),
 )
 

homeassistant/components/justnimbus/strings.json

@@ -3,7 +3,8 @@
     "step": {
       "user": {
         "data": {
-          "client_id": "Client ID"
+          "client_id": "Client ID",
+          "zip_code": "ZIP code"
         }
       }
     },
@@ -13,46 +14,32 @@
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
   },
   "entity": {
     "sensor": {
-      "pump_flow": {
-        "name": "Pump flow"
-      },
-      "drink_flow": {
-        "name": "Drink flow"
-      },
       "pump_pressure": {
         "name": "Pump pressure"
       },
-      "pump_starts": {
+      "pump_type": {
-        "name": "Pump starts"
+        "name": "Pump type"
       },
-      "pump_hours": {
+      "reservoir_capacity": {
-        "name": "Pump hours"
+        "name": "Reservoir capacity"
-      },
-      "reservoir_temperature": {
-        "name": "Reservoir temperature"
       },
       "reservoir_content": {
         "name": "Reservoir content"
       },
-      "total_saved": {
+      "reservoir_temperature": {
-        "name": "Total saved"
+        "name": "Reservoir temperature"
       },
-      "total_replenished": {
+      "water_used": {
-        "name": "Total replenished"
-      },
-      "error_code": {
-        "name": "Error code"
-      },
-      "total_use": {
         "name": "Total use"
       },
-      "reservoir_content_max": {
+      "water_saved": {
-        "name": "Maximum reservoir content"
+        "name": "Total saved"
       }
     }
   }

requirements_all.txt

@@ -1147,7 +1147,7 @@ jellyfin-apiclient-python==1.9.2
 jsonpath==0.82.2
 
 # homeassistant.components.justnimbus
-justnimbus==0.6.0
+justnimbus==0.7.3
 
 # homeassistant.components.kaiterra
 kaiterra-async-client==1.0.0

requirements_test_all.txt

@@ -919,7 +919,7 @@ jellyfin-apiclient-python==1.9.2
 jsonpath==0.82.2
 
 # homeassistant.components.justnimbus
-justnimbus==0.6.0
+justnimbus==0.7.3
 
 # homeassistant.components.kegtron
 kegtron-ble==0.4.0

tests/components/justnimbus/conftest.py

@@ -0,0 +1,8 @@
+"""Reusable fixtures for justnimbus tests."""
+
+from homeassistant.components.justnimbus.const import CONF_ZIP_CODE
+from homeassistant.const import CONF_CLIENT_ID
+
+FIXTURE_OLD_USER_INPUT = {CONF_CLIENT_ID: "test_id"}
+FIXTURE_USER_INPUT = {CONF_CLIENT_ID: "test_id", CONF_ZIP_CODE: "test_zip"}
+FIXTURE_UNIQUE_ID = "test_idtest_zip"

tests/components/justnimbus/test_config_flow.py

@@ -4,12 +4,13 @@ from unittest.mock import patch
 from justnimbus.exceptions import InvalidClientID, JustNimbusError
 import pytest
 
-from homeassistant import config_entries
+from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.justnimbus.const import DOMAIN
-from homeassistant.const import CONF_CLIENT_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from .conftest import FIXTURE_OLD_USER_INPUT, FIXTURE_UNIQUE_ID, FIXTURE_USER_INPUT
+
 from tests.common import MockConfigEntry
 
 
@@ -57,9 +58,7 @@ async def test_form_errors(
     ):
         result2 = await hass.config_entries.flow.async_configure(
             flow_id=result["flow_id"],
-            user_input={
+            user_input=FIXTURE_USER_INPUT,
-                CONF_CLIENT_ID: "test_id",
-            },
         )
 
     assert result2["type"] == FlowResultType.FORM
@@ -73,8 +72,8 @@ async def test_abort_already_configured(hass: HomeAssistant) -> None:
     entry = MockConfigEntry(
         domain=DOMAIN,
         title="JustNimbus",
-        data={CONF_CLIENT_ID: "test_id"},
+        data=FIXTURE_USER_INPUT,
-        unique_id="test_id",
+        unique_id=FIXTURE_UNIQUE_ID,
     )
     entry.add_to_hass(hass)
 
@@ -86,9 +85,7 @@ async def test_abort_already_configured(hass: HomeAssistant) -> None:
 
     result2 = await hass.config_entries.flow.async_configure(
         flow_id=result["flow_id"],
-        user_input={
+        user_input=FIXTURE_USER_INPUT,
-            CONF_CLIENT_ID: "test_id",
-        },
     )
 
     assert result2.get("type") == FlowResultType.ABORT
@@ -103,15 +100,49 @@ async def _set_up_justnimbus(hass: HomeAssistant, flow_id: str) -> None:
     ) as mock_setup_entry:
         result2 = await hass.config_entries.flow.async_configure(
             flow_id=flow_id,
-            user_input={
+            user_input=FIXTURE_USER_INPUT,
-                CONF_CLIENT_ID: "test_id",
-            },
         )
         await hass.async_block_till_done()
 
     assert result2["type"] == FlowResultType.CREATE_ENTRY
     assert result2["title"] == "JustNimbus"
-    assert result2["data"] == {
+    assert result2["data"] == FIXTURE_USER_INPUT
-        CONF_CLIENT_ID: "test_id",
-    }
     assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_reauth_flow(hass: HomeAssistant) -> None:
+    """Test reauth works."""
+    with patch(
+        "homeassistant.components.justnimbus.config_flow.justnimbus.JustNimbusClient.get_data",
+        return_value=False,
+    ):
+        mock_config = MockConfigEntry(
+            domain=DOMAIN, unique_id=FIXTURE_UNIQUE_ID, data=FIXTURE_OLD_USER_INPUT
+        )
+        mock_config.add_to_hass(hass)
+
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={
+                "source": config_entries.SOURCE_REAUTH,
+                "entry_id": mock_config.entry_id,
+            },
+            data=FIXTURE_OLD_USER_INPUT,
+        )
+
+        assert result["type"] == data_entry_flow.FlowResultType.FORM
+        assert result["step_id"] == "user"
+
+    with patch(
+        "homeassistant.components.justnimbus.config_flow.justnimbus.JustNimbusClient.get_data",
+        return_value=True,
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            FIXTURE_USER_INPUT,
+        )
+        await hass.async_block_till_done()
+
+        assert result2["type"] == data_entry_flow.FlowResultType.ABORT
+        assert result2["reason"] == "reauth_successful"
+        assert mock_config.data == FIXTURE_USER_INPUT

tests/components/justnimbus/test_init.py

@@ -0,0 +1,21 @@
+"""Tests for JustNimbus initialization."""
+from homeassistant.components.justnimbus.const import DOMAIN
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from .conftest import FIXTURE_OLD_USER_INPUT, FIXTURE_UNIQUE_ID
+
+from tests.common import MockConfigEntry
+
+
+async def test_config_entry_reauth_at_setup(hass: HomeAssistant) -> None:
+    """Test that setting up with old config results in reauth."""
+    mock_config = MockConfigEntry(
+        domain=DOMAIN, unique_id=FIXTURE_UNIQUE_ID, data=FIXTURE_OLD_USER_INPUT
+    )
+    mock_config.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_config.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config.state is ConfigEntryState.SETUP_ERROR
+    assert any(mock_config.async_get_active_flows(hass, {"reauth"}))