Add switch to control downloads for nzbget (#40673)

* add switch to control downloads for nzbget

* work on switch

* Update test_config_flow.py

* Update test_sensor.py

homeassistant/components/nzbget/__init__.py

@@ -37,7 +37,7 @@ from .coordinator import NZBGetDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = ["sensor"]
+PLATFORMS = ["sensor", "switch"]
 
 CONFIG_SCHEMA = vol.Schema(
     {

homeassistant/components/nzbget/sensor.py

@@ -64,7 +64,7 @@ async def async_setup_entry(
     async_add_entities(sensors)
 
 
-class NZBGetSensor(NZBGetEntity, Entity):
+class NZBGetSensor(NZBGetEntity):
     """Representation of a NZBGet sensor."""
 
     def __init__(

homeassistant/components/nzbget/switch.py

@@ -0,0 +1,72 @@
+"""Support for NZBGet switches."""
+from typing import Callable, List
+
+from homeassistant.components.switch import SwitchEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_NAME
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.typing import HomeAssistantType
+
+from . import NZBGetEntity
+from .const import DATA_COORDINATOR, DOMAIN
+from .coordinator import NZBGetDataUpdateCoordinator
+
+
+async def async_setup_entry(
+    hass: HomeAssistantType,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[List[Entity], bool], None],
+) -> None:
+    """Set up NZBGet sensor based on a config entry."""
+    coordinator: NZBGetDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
+        DATA_COORDINATOR
+    ]
+
+    switches = [
+        NZBGetDownloadSwitch(
+            coordinator,
+            entry.entry_id,
+            entry.data[CONF_NAME],
+        ),
+    ]
+
+    async_add_entities(switches)
+
+
+class NZBGetDownloadSwitch(NZBGetEntity, SwitchEntity):
+    """Representation of a NZBGet download switch."""
+
+    def __init__(
+        self,
+        coordinator: NZBGetDataUpdateCoordinator,
+        entry_id: str,
+        entry_name: str,
+    ):
+        """Initialize a new NZBGet switch."""
+        self._unique_id = f"{entry_id}_download"
+
+        super().__init__(
+            coordinator=coordinator,
+            entry_id=entry_id,
+            name=f"{entry_name} Download",
+        )
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID of the switch."""
+        return self._unique_id
+
+    @property
+    def is_on(self):
+        """Return the state of the switch."""
+        return not self.coordinator.data["status"].get("DownloadPaused", False)
+
+    async def async_turn_on(self, **kwargs) -> None:
+        """Set downloads to enabled."""
+        await self.hass.async_add_executor_job(self.coordinator.nzbget.resumedownload)
+        await self.coordinator.async_request_refresh()
+
+    async def async_turn_off(self, **kwargs) -> None:
+        """Set downloads to paused."""
+        await self.hass.async_add_executor_job(self.coordinator.nzbget.pausedownload)
+        await self.coordinator.async_request_refresh()

tests/components/nzbget/__init__.py

@@ -26,6 +26,8 @@ ENTRY_CONFIG = {
     CONF_VERIFY_SSL: False,
 }
 
+ENTRY_OPTIONS = {CONF_SCAN_INTERVAL: 5}
+
 USER_INPUT = {
     CONF_HOST: "10.10.10.30",
     CONF_NAME: "NZBGet",
@@ -50,12 +52,12 @@ MOCK_VERSION = "21.0"
 MOCK_STATUS = {
     "ArticleCacheMB": 64,
     "AverageDownloadRate": 1250000,
-    "DownloadPaused": 4,
+    "DownloadPaused": False,
     "DownloadRate": 2500000,
     "DownloadedSizeMB": 256,
     "FreeDiskSpaceMB": 1024,
     "PostJobCount": 2,
-    "PostPaused": 4,
+    "PostPaused": False,
     "RemainingSizeMB": 512,
     "UpTimeSec": 600,
 }
@@ -69,17 +71,15 @@ MOCK_HISTORY = [
 async def init_integration(
     hass,
     *,
-    status: dict = MOCK_STATUS,
-    history: dict = MOCK_HISTORY,
-    version: str = MOCK_VERSION,
+    data: dict = ENTRY_CONFIG,
+    options: dict = ENTRY_OPTIONS,
 ) -> MockConfigEntry:
     """Set up the NZBGet integration in Home Assistant."""
-    entry = MockConfigEntry(domain=DOMAIN, data=ENTRY_CONFIG)
+    entry = MockConfigEntry(domain=DOMAIN, data=data, options=options)
     entry.add_to_hass(hass)
 
-    with _patch_version(version), _patch_status(status), _patch_history(history):
-        await hass.config_entries.async_setup(entry.entry_id)
-        await hass.async_block_till_done()
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
 
     return entry
 

tests/components/nzbget/conftest.py

@@ -0,0 +1,21 @@
+"""Define fixtures available for all tests."""
+from pytest import fixture
+
+from . import MOCK_HISTORY, MOCK_STATUS, MOCK_VERSION
+
+from tests.async_mock import MagicMock, patch
+
+
+@fixture
+def nzbget_api(hass):
+    """Mock NZBGetApi for easier testing."""
+    with patch("homeassistant.components.nzbget.coordinator.NZBGetAPI") as mock_api:
+        instance = mock_api.return_value
+
+        instance.history = MagicMock(return_value=list(MOCK_HISTORY))
+        instance.pausedownload = MagicMock(return_value=True)
+        instance.resumedownload = MagicMock(return_value=True)
+        instance.status = MagicMock(return_value=MOCK_STATUS.copy())
+        instance.version = MagicMock(return_value=MOCK_VERSION)
+
+        yield mock_api

tests/components/nzbget/test_config_flow.py

@@ -132,7 +132,7 @@ async def test_user_form_single_instance_allowed(hass):
     assert result["reason"] == "single_instance_allowed"
 
 
-async def test_options_flow(hass):
+async def test_options_flow(hass, nzbget_api):
     """Test updating options."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -141,16 +141,22 @@ async def test_options_flow(hass):
     )
     entry.add_to_hass(hass)
 
+    with patch("homeassistant.components.nzbget.PLATFORMS", []):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
     assert entry.options[CONF_SCAN_INTERVAL] == 5
 
     result = await hass.config_entries.options.async_init(entry.entry_id)
     assert result["type"] == RESULT_TYPE_FORM
     assert result["step_id"] == "init"
 
-    result = await hass.config_entries.options.async_configure(
-        result["flow_id"],
-        user_input={CONF_SCAN_INTERVAL: 15},
-    )
+    with _patch_async_setup(), _patch_async_setup_entry():
+        result = await hass.config_entries.options.async_configure(
+            result["flow_id"],
+            user_input={CONF_SCAN_INTERVAL: 15},
+        )
+        await hass.async_block_till_done()
 
     assert result["type"] == RESULT_TYPE_CREATE_ENTRY
     assert result["data"][CONF_SCAN_INTERVAL] == 15

tests/components/nzbget/test_init.py

@@ -38,7 +38,7 @@ async def test_import_from_yaml(hass) -> None:
     assert entries[0].data[CONF_PORT] == 6789
 
 
-async def test_unload_entry(hass):
+async def test_unload_entry(hass, nzbget_api):
     """Test successful unload of entry."""
     entry = await init_integration(hass)
 

tests/components/nzbget/test_sensor.py

@@ -14,7 +14,7 @@ from . import init_integration
 from tests.async_mock import patch
 
 
-async def test_sensors(hass) -> None:
+async def test_sensors(hass, nzbget_api) -> None:
     """Test the creation and values of the sensors."""
     now = dt_util.utcnow().replace(microsecond=0)
     with patch("homeassistant.util.dt.utcnow", return_value=now):
@@ -32,12 +32,12 @@ async def test_sensors(hass) -> None:
             DATA_RATE_MEGABYTES_PER_SECOND,
             None,
         ),
-        "download_paused": ("DownloadPaused", "4", None, None),
+        "download_paused": ("DownloadPaused", "False", None, None),
         "speed": ("DownloadRate", "2.38", DATA_RATE_MEGABYTES_PER_SECOND, None),
         "size": ("DownloadedSizeMB", "256", DATA_MEGABYTES, None),
         "disk_free": ("FreeDiskSpaceMB", "1024", DATA_MEGABYTES, None),
         "post_processing_jobs": ("PostJobCount", "2", "Jobs", None),
-        "post_processing_paused": ("PostPaused", "4", None, None),
+        "post_processing_paused": ("PostPaused", "False", None, None),
         "queue_size": ("RemainingSizeMB", "512", DATA_MEGABYTES, None),
         "uptime": ("UpTimeSec", uptime.isoformat(), None, DEVICE_CLASS_TIMESTAMP),
     }

tests/components/nzbget/test_switch.py

@@ -0,0 +1,64 @@
+"""Test the NZBGet switches."""
+from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+)
+
+from . import init_integration
+
+
+async def test_download_switch(hass, nzbget_api) -> None:
+    """Test the creation and values of the download switch."""
+    instance = nzbget_api.return_value
+
+    entry = await init_integration(hass)
+    assert entry
+
+    registry = await hass.helpers.entity_registry.async_get_registry()
+    entity_id = "switch.nzbgettest_download"
+    entity_entry = registry.async_get(entity_id)
+    assert entity_entry
+    assert entity_entry.unique_id == f"{entry.entry_id}_download"
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_ON
+
+    # test download paused
+    instance.status.return_value["DownloadPaused"] = True
+
+    await hass.helpers.entity_component.async_update_entity(entity_id)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_OFF
+
+
+async def test_download_switch_services(hass, nzbget_api) -> None:
+    """Test download switch services."""
+    instance = nzbget_api.return_value
+
+    entry = await init_integration(hass)
+    entity_id = "switch.nzbgettest_download"
+    assert entry
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    instance.pausedownload.assert_called_once()
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    instance.resumedownload.assert_called_once()