From fde7ddfa7148fef7e6373d636edbff4bce5355b8 Mon Sep 17 00:00:00 2001
From: Brett Adams <Bre77@users.noreply.github.com>
Date: Fri, 21 Jun 2024 19:30:57 +1000
Subject: [PATCH] Fix charge behavior in Tessie (#119546)

---
 homeassistant/components/tessie/entity.py    |  2 +-
 homeassistant/components/tessie/switch.py    | 43 ++++++++++++++------
 tests/components/tessie/fixtures/online.json |  2 +-
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/tessie/entity.py b/homeassistant/components/tessie/entity.py
index e11a99348ed..35d41af32f2 100644
--- a/homeassistant/components/tessie/entity.py
+++ b/homeassistant/components/tessie/entity.py
@@ -46,7 +46,7 @@ class TessieEntity(CoordinatorEntity[TessieStateUpdateCoordinator]):
     @property
     def _value(self) -> Any:
         """Return value from coordinator data."""
-        return self.coordinator.data[self.key]
+        return self.coordinator.data.get(self.key)
 
     def get(self, key: str | None = None, default: Any | None = None) -> Any:
         """Return a specific value from coordinator data."""
diff --git a/homeassistant/components/tessie/switch.py b/homeassistant/components/tessie/switch.py
index 191d4f3ff5c..2f3902b3bd3 100644
--- a/homeassistant/components/tessie/switch.py
+++ b/homeassistant/components/tessie/switch.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
+from itertools import chain
 from typing import Any
 
 from tessie_api import (
@@ -41,11 +42,6 @@ class TessieSwitchEntityDescription(SwitchEntityDescription):
 
 
 DESCRIPTIONS: tuple[TessieSwitchEntityDescription, ...] = (
-    TessieSwitchEntityDescription(
-        key="charge_state_charge_enable_request",
-        on_func=lambda: start_charging,
-        off_func=lambda: stop_charging,
-    ),
     TessieSwitchEntityDescription(
         key="climate_state_defrost_mode",
         on_func=lambda: start_defrost,
@@ -68,6 +64,12 @@ DESCRIPTIONS: tuple[TessieSwitchEntityDescription, ...] = (
     ),
 )
 
+CHARGE_DESCRIPTION: TessieSwitchEntityDescription = TessieSwitchEntityDescription(
+    key="charge_state_charge_enable_request",
+    on_func=lambda: start_charging,
+    off_func=lambda: stop_charging,
+)
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -75,15 +77,20 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Tessie Switch platform from a config entry."""
-    data = entry.runtime_data
 
     async_add_entities(
-        [
-            TessieSwitchEntity(vehicle, description)
-            for vehicle in data.vehicles
-            for description in DESCRIPTIONS
-            if description.key in vehicle.data
-        ]
+        chain(
+            (
+                TessieSwitchEntity(vehicle, description)
+                for vehicle in entry.runtime_data.vehicles
+                for description in DESCRIPTIONS
+                if description.key in vehicle.data
+            ),
+            (
+                TessieChargeSwitchEntity(vehicle, CHARGE_DESCRIPTION)
+                for vehicle in entry.runtime_data.vehicles
+            ),
+        )
     )
 
 
@@ -116,3 +123,15 @@ class TessieSwitchEntity(TessieEntity, SwitchEntity):
         """Turn off the Switch."""
         await self.run(self.entity_description.off_func())
         self.set((self.entity_description.key, False))
+
+
+class TessieChargeSwitchEntity(TessieSwitchEntity):
+    """Entity class for Tessie charge switch."""
+
+    @property
+    def is_on(self) -> bool:
+        """Return the state of the Switch."""
+
+        if (charge := self.get("charge_state_user_charge_enable_request")) is not None:
+            return charge
+        return self._value
diff --git a/tests/components/tessie/fixtures/online.json b/tests/components/tessie/fixtures/online.json
index 863e9bca783..ed49b4bfd75 100644
--- a/tests/components/tessie/fixtures/online.json
+++ b/tests/components/tessie/fixtures/online.json
@@ -68,7 +68,7 @@
     "timestamp": 1701139037461,
     "trip_charging": false,
     "usable_battery_level": 75,
-    "user_charge_enable_request": null
+    "user_charge_enable_request": true
   },
   "climate_state": {
     "allow_cabin_overheat_protection": true,