From 7840db0598318061ea19e015c1df9e633f0cc34c Mon Sep 17 00:00:00 2001
From: Matt Zimmerman <mdz@users.noreply.github.com>
Date: Tue, 9 Mar 2021 16:26:54 -0800
Subject: [PATCH] Restore switch.turn_on and switch.turn_off functionality for
 SmartTub pumps (#47586)

Revert "Remove turn_on and turn_off from SmartTub pump switches (#47184)"

This reverts commit bab66a5cb968bcc7c6ef7787c9dc645885569429.
---
 homeassistant/components/smarttub/switch.py | 14 ++++++++++++++
 tests/components/smarttub/test_switch.py    | 21 +++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/homeassistant/components/smarttub/switch.py b/homeassistant/components/smarttub/switch.py
index d2891932546..26239df9dff 100644
--- a/homeassistant/components/smarttub/switch.py
+++ b/homeassistant/components/smarttub/switch.py
@@ -61,6 +61,20 @@ class SmartTubPump(SmartTubEntity, SwitchEntity):
         """Return True if the pump is on."""
         return self.pump.state != SpaPump.PumpState.OFF
 
+    async def async_turn_on(self, **kwargs) -> None:
+        """Turn the pump on."""
+
+        # the API only supports toggling
+        if not self.is_on:
+            await self.async_toggle()
+
+    async def async_turn_off(self, **kwargs) -> None:
+        """Turn the pump off."""
+
+        # the API only supports toggling
+        if self.is_on:
+            await self.async_toggle()
+
     async def async_toggle(self, **kwargs) -> None:
         """Toggle the pump on or off."""
         async with async_timeout.timeout(API_TIMEOUT):
diff --git a/tests/components/smarttub/test_switch.py b/tests/components/smarttub/test_switch.py
index 89e0ec03b23..81b53604065 100644
--- a/tests/components/smarttub/test_switch.py
+++ b/tests/components/smarttub/test_switch.py
@@ -2,6 +2,8 @@
 
 import pytest
 
+from homeassistant.const import STATE_OFF, STATE_ON
+
 
 @pytest.mark.parametrize(
     "pump_id,entity_suffix,pump_state",
@@ -28,3 +30,22 @@ async def test_pumps(spa, setup_entry, hass, pump_id, pump_state, entity_suffix)
         blocking=True,
     )
     pump.toggle.assert_called()
+
+    if state.state == STATE_OFF:
+        await hass.services.async_call(
+            "switch",
+            "turn_on",
+            {"entity_id": entity_id},
+            blocking=True,
+        )
+        pump.toggle.assert_called()
+    else:
+        assert state.state == STATE_ON
+
+        await hass.services.async_call(
+            "switch",
+            "turn_off",
+            {"entity_id": entity_id},
+            blocking=True,
+        )
+        pump.toggle.assert_called()