Add price cap support to Ohme (#140537)

* Add price cap support

* Change service input to box mode

* Add icon for set_price_cap service

* Improve test coverage

* Change ohme service description wording

homeassistant/components/ohme/icons.json

@@ -54,6 +54,9 @@
         "state": {
           "off": "mdi:sleep-off"
         }
+      },
+      "price_cap": {
+        "default": "mdi:car-speed-limiter"
       }
     },
     "time": {
@@ -65,6 +68,9 @@
   "services": {
     "list_charge_slots": {
       "service": "mdi:clock-start"
+    },
+    "set_price_cap": {
+      "service": "mdi:car-speed-limiter"
     }
   }
 }

homeassistant/components/ohme/services.py

@@ -17,9 +17,11 @@ from homeassistant.helpers import selector
 
 from .const import DOMAIN
 
-SERVICE_LIST_CHARGE_SLOTS = "list_charge_slots"
 ATTR_CONFIG_ENTRY: Final = "config_entry"
-SERVICE_SCHEMA: Final = vol.Schema(
+ATTR_PRICE_CAP: Final = "price_cap"
+
+SERVICE_LIST_CHARGE_SLOTS = "list_charge_slots"
+SERVICE_LIST_CHARGE_SLOTS_SCHEMA: Final = vol.Schema(
     {
         vol.Required(ATTR_CONFIG_ENTRY): selector.ConfigEntrySelector(
             {
@@ -29,6 +31,18 @@ SERVICE_SCHEMA: Final = vol.Schema(
     }
 )
 
+SERVICE_SET_PRICE_CAP = "set_price_cap"
+SERVICE_SET_PRICE_CAP_SCHEMA: Final = vol.Schema(
+    {
+        vol.Required(ATTR_CONFIG_ENTRY): selector.ConfigEntrySelector(
+            {
+                "integration": DOMAIN,
+            }
+        ),
+        vol.Required(ATTR_PRICE_CAP): vol.Coerce(float),
+    }
+)
+
 
 def __get_client(call: ServiceCall) -> OhmeApiClient:
     """Get the client from the config entry."""
@@ -66,10 +80,26 @@ def async_setup_services(hass: HomeAssistant) -> None:
 
         return {"slots": client.slots}
 
+    async def set_price_cap(
+        service_call: ServiceCall,
+    ) -> None:
+        """List of charge slots."""
+        client = __get_client(service_call)
+        price_cap = service_call.data[ATTR_PRICE_CAP]
+        await client.async_change_price_cap(cap=price_cap)
+
     hass.services.async_register(
         DOMAIN,
         SERVICE_LIST_CHARGE_SLOTS,
         list_charge_slots,
-        schema=SERVICE_SCHEMA,
+        schema=SERVICE_LIST_CHARGE_SLOTS_SCHEMA,
         supports_response=SupportsResponse.ONLY,
     )
+
+    hass.services.async_register(
+        DOMAIN,
+        SERVICE_SET_PRICE_CAP,
+        set_price_cap,
+        schema=SERVICE_SET_PRICE_CAP_SCHEMA,
+        supports_response=SupportsResponse.NONE,
+    )

homeassistant/components/ohme/services.yaml

@@ -5,3 +5,16 @@ list_charge_slots:
       selector:
         config_entry:
           integration: ohme
+set_price_cap:
+  fields:
+    config_entry:
+      required: true
+      selector:
+        config_entry:
+          integration: ohme
+    price_cap:
+      required: true
+      selector:
+        number:
+          min: 0
+          mode: box

homeassistant/components/ohme/strings.json

@@ -42,6 +42,20 @@
           "description": "The Ohme config entry for which to return charge slots."
         }
       }
+    },
+    "set_price_cap": {
+      "name": "Set price cap",
+      "description": "Prevents charging when the electricity price exceeds a defined threshold.",
+      "fields": {
+        "config_entry": {
+          "name": "Ohme account",
+          "description": "The Ohme config entry for which to return charge slots."
+        },
+        "price_cap": {
+          "name": "Price cap",
+          "description": "Threshold in 1/100ths of your local currency."
+        }
+      }
     }
   },
   "entity": {
@@ -102,6 +116,9 @@
       },
       "sleep_when_inactive": {
         "name": "Sleep when inactive"
+      },
+      "price_cap": {
+        "name": "Price cap"
       }
     },
     "time": {

homeassistant/components/ohme/switch.py

@@ -1,9 +1,10 @@
 """Platform for switch."""
 
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any
 
-from ohme import ApiException
+from ohme import ApiException, OhmeApiClient
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.const import EntityCategory
@@ -19,28 +20,37 @@ PARALLEL_UPDATES = 1
 
 
 @dataclass(frozen=True, kw_only=True)
-class OhmeSwitchDescription(OhmeEntityDescription, SwitchEntityDescription):
-    """Class describing Ohme switch entities."""
+class OhmeConfigSwitchDescription(OhmeEntityDescription, SwitchEntityDescription):
+    """Class describing Ohme configuration switch entities."""
 
     configuration_key: str
 
 
-SWITCH_DEVICE_INFO = [
-    OhmeSwitchDescription(
+@dataclass(frozen=True, kw_only=True)
+class OhmeSwitchDescription(OhmeEntityDescription, SwitchEntityDescription):
+    """Class describing basic Ohme switch entities."""
+
+    is_on_fn: Callable[[OhmeApiClient], bool]
+    off_fn: Callable[[OhmeApiClient], Awaitable]
+    on_fn: Callable[[OhmeApiClient], Awaitable]
+
+
+SWITCH_CONFIG = [
+    OhmeConfigSwitchDescription(
         key="lock_buttons",
         translation_key="lock_buttons",
         entity_category=EntityCategory.CONFIG,
         is_supported_fn=lambda client: client.is_capable("buttonsLockable"),
         configuration_key="buttonsLocked",
     ),
-    OhmeSwitchDescription(
+    OhmeConfigSwitchDescription(
         key="require_approval",
         translation_key="require_approval",
         entity_category=EntityCategory.CONFIG,
         is_supported_fn=lambda client: client.is_capable("pluginsRequireApprovalMode"),
         configuration_key="pluginsRequireApproval",
     ),
-    OhmeSwitchDescription(
+    OhmeConfigSwitchDescription(
         key="sleep_when_inactive",
         translation_key="sleep_when_inactive",
         entity_category=EntityCategory.CONFIG,
@@ -49,6 +59,17 @@ SWITCH_DEVICE_INFO = [
     ),
 ]
 
+SWITCH_DESCRIPTION = [
+    OhmeSwitchDescription(
+        key="price_cap",
+        translation_key="price_cap",
+        is_supported_fn=lambda client: client.cap_available,
+        is_on_fn=lambda client: client.cap_enabled,
+        on_fn=lambda client: client.async_change_price_cap(True),
+        off_fn=lambda client: client.async_change_price_cap(False),
+    ),
+]
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -56,15 +77,17 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up switches."""
-    coordinators = config_entry.runtime_data
-    coordinator_map = [
-        (SWITCH_DEVICE_INFO, coordinators.device_info_coordinator),
-    ]
+    coordinator = config_entry.runtime_data.device_info_coordinator
+
+    async_add_entities(
+        OhmeConfigSwitch(coordinator, description)
+        for description in SWITCH_CONFIG
+        if description.is_supported_fn(coordinator.client)
+    )
 
     async_add_entities(
         OhmeSwitch(coordinator, description)
-        for entities, coordinator in coordinator_map
-        for description in entities
+        for description in SWITCH_DESCRIPTION
         if description.is_supported_fn(coordinator.client)
     )
 
@@ -74,6 +97,27 @@ class OhmeSwitch(OhmeEntity, SwitchEntity):
 
     entity_description: OhmeSwitchDescription
 
+    @property
+    def is_on(self) -> bool:
+        """Return True if entity is on."""
+        return self.entity_description.is_on_fn(self.coordinator.client)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the switch off."""
+        await self.entity_description.off_fn(self.coordinator.client)
+        await self.coordinator.async_request_refresh()
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the switch on."""
+        await self.entity_description.on_fn(self.coordinator.client)
+        await self.coordinator.async_request_refresh()
+
+
+class OhmeConfigSwitch(OhmeEntity, SwitchEntity):
+    """Configuration switch for Ohme."""
+
+    entity_description: OhmeConfigSwitchDescription
+
     @property
     def is_on(self) -> bool:
         """Return the entity value to represent the entity state."""

tests/components/ohme/conftest.py

@@ -60,6 +60,8 @@ def mock_client():
         client.preconditioning = 15
         client.serial = "chargerid"
         client.ct_connected = True
+        client.cap_available = True
+        client.cap_enabled = True
         client.energy = 1000
         client.device_info = {
             "name": "Ohme Home Pro",

tests/components/ohme/snapshots/test_switch.ambr

@@ -46,6 +46,53 @@
     'state': 'on',
   })
 # ---
+# name: test_switches[switch.ohme_home_pro_price_cap-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'config_subentry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'switch',
+    'entity_category': None,
+    'entity_id': 'switch.ohme_home_pro_price_cap',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Price cap',
+    'platform': 'ohme',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'price_cap',
+    'unique_id': 'chargerid_price_cap',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_switches[switch.ohme_home_pro_price_cap-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Ohme Home Pro Price cap',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.ohme_home_pro_price_cap',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
 # name: test_switches[switch.ohme_home_pro_require_approval-entry]
   EntityRegistryEntrySnapshot({
     'aliases': set({

tests/components/ohme/test_services.py

@@ -1,6 +1,6 @@
 """Tests for services."""
 
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -8,6 +8,7 @@ from syrupy.assertion import SnapshotAssertion
 from homeassistant.components.ohme.const import DOMAIN
 from homeassistant.components.ohme.services import (
     ATTR_CONFIG_ENTRY,
+    ATTR_PRICE_CAP,
     SERVICE_LIST_CHARGE_SLOTS,
 )
 from homeassistant.core import HomeAssistant
@@ -47,6 +48,30 @@ async def test_list_charge_slots(
     )
 
 
+async def test_set_price_cap(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_client: MagicMock,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test set price cap service."""
+
+    await setup_integration(hass, mock_config_entry)
+    mock_client.async_change_price_cap = AsyncMock()
+
+    await hass.services.async_call(
+        DOMAIN,
+        "set_price_cap",
+        {
+            ATTR_CONFIG_ENTRY: mock_config_entry.entry_id,
+            ATTR_PRICE_CAP: 10.0,
+        },
+        blocking=True,
+    )
+
+    mock_client.async_change_price_cap.assert_called_once_with(cap=10.0)
+
+
 async def test_list_charge_slots_exception(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,

tests/components/ohme/test_switch.py

@@ -1,6 +1,6 @@
 """Tests for switches."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from syrupy import SnapshotAssertion
 
@@ -32,7 +32,49 @@ async def test_switches(
     await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
-async def test_switch_on(
+async def test_cap_switch_on(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_client: MagicMock,
+) -> None:
+    """Test the switch turn_on action."""
+    await setup_integration(hass, mock_config_entry)
+    mock_client.async_change_price_cap = AsyncMock()
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {
+            ATTR_ENTITY_ID: "switch.ohme_home_pro_price_cap",
+        },
+        blocking=True,
+    )
+
+    mock_client.async_change_price_cap.assert_called_once_with(True)
+
+
+async def test_cap_switch_off(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_client: MagicMock,
+) -> None:
+    """Test the switch turn_off action."""
+    await setup_integration(hass, mock_config_entry)
+    mock_client.async_change_price_cap = AsyncMock()
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {
+            ATTR_ENTITY_ID: "switch.ohme_home_pro_price_cap",
+        },
+        blocking=True,
+    )
+
+    mock_client.async_change_price_cap.assert_called_once_with(False)
+
+
+async def test_config_switch_on(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     mock_client: MagicMock,
@@ -52,7 +94,7 @@ async def test_switch_on(
     assert len(mock_client.async_set_configuration_value.mock_calls) == 1
 
 
-async def test_switch_off(
+async def test_config_switch_off(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     mock_client: MagicMock,