From 1b42b32ac10d46e1a13cae80452288abf4b6a670 Mon Sep 17 00:00:00 2001
From: Thomas55555 <59625598+Thomas55555@users.noreply.github.com>
Date: Fri, 5 Jul 2024 09:41:21 +0200
Subject: [PATCH] Fix work area sensor in Husqvarna Automower (#121228)

---
 .../components/husqvarna_automower/sensor.py  | 11 +++++--
 .../husqvarna_automower/strings.json          |  3 +-
 .../snapshots/test_sensor.ambr                |  2 ++
 .../husqvarna_automower/test_sensor.py        | 32 +++++++++++++++++++
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/husqvarna_automower/sensor.py b/homeassistant/components/husqvarna_automower/sensor.py
index 146ef17a6e4..2c8d369ea3a 100644
--- a/homeassistant/components/husqvarna_automower/sensor.py
+++ b/homeassistant/components/husqvarna_automower/sensor.py
@@ -184,6 +184,8 @@ RESTRICTED_REASONS: list = [
     RestrictedReasons.WEEK_SCHEDULE.lower(),
 ]
 
+STATE_NO_WORK_AREA_ACTIVE = "no_work_area_active"
+
 
 @callback
 def _get_work_area_names(data: MowerAttributes) -> list[str]:
@@ -191,16 +193,21 @@ def _get_work_area_names(data: MowerAttributes) -> list[str]:
     if TYPE_CHECKING:
         # Sensor does not get created if it is None
         assert data.work_areas is not None
-    return [data.work_areas[work_area_id].name for work_area_id in data.work_areas]
+    work_area_list = [
+        data.work_areas[work_area_id].name for work_area_id in data.work_areas
+    ]
+    work_area_list.append(STATE_NO_WORK_AREA_ACTIVE)
+    return work_area_list
 
 
 @callback
 def _get_current_work_area_name(data: MowerAttributes) -> str:
     """Return the name of the current work area."""
+    if data.mower.work_area_id is None:
+        return STATE_NO_WORK_AREA_ACTIVE
     if TYPE_CHECKING:
         # Sensor does not get created if values are None
         assert data.work_areas is not None
-        assert data.mower.work_area_id is not None
     return data.work_areas[data.mower.work_area_id].name
 
 
diff --git a/homeassistant/components/husqvarna_automower/strings.json b/homeassistant/components/husqvarna_automower/strings.json
index 6cb1c17421a..be17cc25e32 100644
--- a/homeassistant/components/husqvarna_automower/strings.json
+++ b/homeassistant/components/husqvarna_automower/strings.json
@@ -252,7 +252,8 @@
       "work_area": {
         "name": "Work area",
         "state": {
-          "my_lawn": "My lawn"
+          "my_lawn": "My lawn",
+          "no_work_area_active": "No work area active"
         }
       }
     },
diff --git a/tests/components/husqvarna_automower/snapshots/test_sensor.ambr b/tests/components/husqvarna_automower/snapshots/test_sensor.ambr
index 0b0d76620d3..935303e48fb 100644
--- a/tests/components/husqvarna_automower/snapshots/test_sensor.ambr
+++ b/tests/components/husqvarna_automower/snapshots/test_sensor.ambr
@@ -1059,6 +1059,7 @@
         'Front lawn',
         'Back lawn',
         'my_lawn',
+        'no_work_area_active',
       ]),
     }),
     'config_entry_id': <ANY>,
@@ -1097,6 +1098,7 @@
         'Front lawn',
         'Back lawn',
         'my_lawn',
+        'no_work_area_active',
       ]),
     }),
     'context': <ANY>,
diff --git a/tests/components/husqvarna_automower/test_sensor.py b/tests/components/husqvarna_automower/test_sensor.py
index 8f30a3dcb04..314bcaaa00c 100644
--- a/tests/components/husqvarna_automower/test_sensor.py
+++ b/tests/components/husqvarna_automower/test_sensor.py
@@ -87,6 +87,38 @@ async def test_next_start_sensor(
     assert state.state == STATE_UNKNOWN
 
 
+async def test_work_area_sensor(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test the work area sensor."""
+    await setup_integration(hass, mock_config_entry)
+    state = hass.states.get("sensor.test_mower_1_work_area")
+    assert state is not None
+    assert state.state == "Front lawn"
+
+    values = mower_list_to_dictionary_dataclass(
+        load_json_value_fixture("mower.json", DOMAIN)
+    )
+    values[TEST_MOWER_ID].mower.work_area_id = None
+    mock_automower_client.get_status.return_value = values
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.test_mower_1_work_area")
+    assert state.state == "no_work_area_active"
+
+    values[TEST_MOWER_ID].mower.work_area_id = 0
+    mock_automower_client.get_status.return_value = values
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.test_mower_1_work_area")
+    assert state.state == "my_lawn"
+
+
 @pytest.mark.parametrize(
     ("sensor_to_test"),
     [