From d50b5cebeec776aecb5f3236d9c4090c69003100 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 16 Aug 2022 17:10:11 +0200
Subject: [PATCH] Various improvement for JustNimbus (#76858)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
---
 homeassistant/components/justnimbus/entity.py | 10 ++----
 homeassistant/components/justnimbus/sensor.py |  8 ++---
 .../components/justnimbus/test_config_flow.py | 36 ++++++++++++++++++-
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/homeassistant/components/justnimbus/entity.py b/homeassistant/components/justnimbus/entity.py
index f9ea5ba1151..cfd261f1bc0 100644
--- a/homeassistant/components/justnimbus/entity.py
+++ b/homeassistant/components/justnimbus/entity.py
@@ -1,19 +1,15 @@
 """Base Entity for JustNimbus sensors."""
 from __future__ import annotations
 
-import justnimbus
-
-from homeassistant.components.sensor import SensorEntity
-from homeassistant.helpers import update_coordinator
 from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import JustNimbusCoordinator
 from .const import DOMAIN
+from .coordinator import JustNimbusCoordinator
 
 
 class JustNimbusEntity(
-    update_coordinator.CoordinatorEntity[justnimbus.JustNimbusModel],
-    SensorEntity,
+    CoordinatorEntity[JustNimbusCoordinator],
 ):
     """Defines a base JustNimbus entity."""
 
diff --git a/homeassistant/components/justnimbus/sensor.py b/homeassistant/components/justnimbus/sensor.py
index 6041f84e25a..73a68ac9139 100644
--- a/homeassistant/components/justnimbus/sensor.py
+++ b/homeassistant/components/justnimbus/sensor.py
@@ -7,6 +7,7 @@ from typing import Any
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
+    SensorEntity,
     SensorEntityDescription,
     SensorStateClass,
 )
@@ -15,6 +16,7 @@ from homeassistant.const import (
     CONF_CLIENT_ID,
     PRESSURE_BAR,
     TEMP_CELSIUS,
+    TIME_HOURS,
     VOLUME_LITERS,
 )
 from homeassistant.core import HomeAssistant
@@ -82,6 +84,7 @@ SENSOR_TYPES = (
         name="Pump hours",
         icon="mdi:clock",
         device_class=SensorDeviceClass.DURATION,
+        native_unit_of_measurement=TIME_HOURS,
         state_class=SensorStateClass.MEASUREMENT,
         entity_category=EntityCategory.DIAGNOSTIC,
         value_fn=lambda coordinator: coordinator.data.pump_hours,
@@ -127,7 +130,6 @@ SENSOR_TYPES = (
         name="Error code",
         icon="mdi:bug",
         entity_registry_enabled_default=False,
-        native_unit_of_measurement="",
         entity_category=EntityCategory.DIAGNOSTIC,
         value_fn=lambda coordinator: coordinator.data.error_code,
     ),
@@ -167,9 +169,7 @@ async def async_setup_entry(
     )
 
 
-class JustNimbusSensor(
-    JustNimbusEntity,
-):
+class JustNimbusSensor(JustNimbusEntity, SensorEntity):
     """Implementation of the JustNimbus sensor."""
 
     def __init__(
diff --git a/tests/components/justnimbus/test_config_flow.py b/tests/components/justnimbus/test_config_flow.py
index d2cfb64d4c7..1d3565c21bd 100644
--- a/tests/components/justnimbus/test_config_flow.py
+++ b/tests/components/justnimbus/test_config_flow.py
@@ -10,6 +10,8 @@ from homeassistant.const import CONF_CLIENT_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from tests.common import MockConfigEntry
+
 
 async def test_form(hass: HomeAssistant) -> None:
     """Test we get the form."""
@@ -30,9 +32,13 @@ async def test_form(hass: HomeAssistant) -> None:
             {"base": "invalid_auth"},
         ),
         (
-            JustNimbusError(),
+            JustNimbusError,
             {"base": "cannot_connect"},
         ),
+        (
+            RuntimeError,
+            {"base": "unknown"},
+        ),
     ),
 )
 async def test_form_errors(
@@ -62,6 +68,34 @@ async def test_form_errors(
     await _set_up_justnimbus(hass=hass, flow_id=result["flow_id"])
 
 
+async def test_abort_already_configured(hass: HomeAssistant) -> None:
+    """Test we abort when the device is already configured."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        title="JustNimbus",
+        data={CONF_CLIENT_ID: "test_id"},
+        unique_id="test_id",
+    )
+    entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result.get("type") == FlowResultType.FORM
+    assert result.get("errors") is None
+    assert "flow_id" in result
+
+    result2 = await hass.config_entries.flow.async_configure(
+        flow_id=result["flow_id"],
+        user_input={
+            CONF_CLIENT_ID: "test_id",
+        },
+    )
+
+    assert result2.get("type") == FlowResultType.ABORT
+    assert result2.get("reason") == "already_configured"
+
+
 async def _set_up_justnimbus(hass: HomeAssistant, flow_id: str) -> None:
     """Reusable successful setup of JustNimbus sensor."""
     with patch("justnimbus.JustNimbusClient.get_data"), patch(