Improve Snoo testing (#139302)

* improve snoo testing

* change to asyncMock method of testing

* Apply suggestions from code review

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* address comments

* Apply suggestions from code review

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* adress comments

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

tests/components/snoo/__init__.py

@@ -1,5 +1,11 @@
 """Tests for the Happiest Baby Snoo integration."""
 
+from collections.abc import Awaitable, Callable
+from unittest.mock import AsyncMock
+
+import pytest
+from python_snoo.containers import SnooData
+
 from homeassistant.components.snoo.const import DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
@@ -36,3 +42,13 @@ async def async_init_integration(hass: HomeAssistant) -> ConfigEntry:
     await hass.async_block_till_done()
 
     return entry
+
+
+def find_update_callback(
+    mock: AsyncMock, serial_number: str
+) -> Callable[[SnooData], Awaitable[None]]:
+    """Find the update callback for a specific identifier."""
+    for call in mock.subscribe.call_args_list:
+        if call[0][0].serialNumber == serial_number:
+            return call[0][1]
+    pytest.fail(f"Callback for identifier {serial_number} not found")

tests/components/snoo/conftest.py

@@ -5,9 +5,8 @@ from unittest.mock import AsyncMock, patch
 
 import pytest
 from python_snoo.containers import SnooDevice
-from python_snoo.snoo import Snoo
 
-from .const import MOCK_AMAZON_AUTH, MOCK_SNOO_AUTH, MOCK_SNOO_DEVICES
+from .const import MOCK_SNOO_DEVICES, MOCKED_AUTH
 
 
 @pytest.fixture
@@ -19,55 +18,14 @@ def mock_setup_entry() -> Generator[AsyncMock]:
         yield mock_setup_entry
 
 
-class MockedSnoo(Snoo):
-    """Mock the Snoo object."""
-
-    def __init__(self, email, password, clientsession) -> None:
-        """Set up a Mocked Snoo."""
-        super().__init__(email, password, clientsession)
-        self.auth_error = None
-
-    async def subscribe(self, device: SnooDevice, function):
-        """Mock the subscribe function."""
-        return AsyncMock()
-
-    async def send_command(self, command: str, device: SnooDevice, **kwargs):
-        """Mock the send command function."""
-        return AsyncMock()
-
-    async def authorize(self):
-        """Do normal auth flow unless error is patched."""
-        if self.auth_error:
-            raise self.auth_error
-        return await super().authorize()
-
-    def set_auth_error(self, error: Exception | None):
-        """Set an error for authentication."""
-        self.auth_error = error
-
-    async def auth_amazon(self):
-        """Mock the amazon auth."""
-        return MOCK_AMAZON_AUTH
-
-    async def auth_snoo(self, id_token):
-        """Mock the snoo auth."""
-        return MOCK_SNOO_AUTH
-
-    async def schedule_reauthorization(self, snoo_expiry: int):
-        """Mock scheduling reauth."""
-        return AsyncMock()
-
-    async def get_devices(self) -> list[SnooDevice]:
-        """Move getting devices."""
-        return [SnooDevice.from_dict(dev) for dev in MOCK_SNOO_DEVICES]
-
-
 @pytest.fixture(name="bypass_api")
-def bypass_api() -> MockedSnoo:
+def bypass_api() -> Generator[AsyncMock]:
     """Bypass the Snoo api."""
-    api = MockedSnoo("email", "password", AsyncMock())
     with (
-        patch("homeassistant.components.snoo.Snoo", return_value=api),
-        patch("homeassistant.components.snoo.config_flow.Snoo", return_value=api),
+        patch("homeassistant.components.snoo.Snoo", autospec=True) as mock_client,
+        patch("homeassistant.components.snoo.config_flow.Snoo", new=mock_client),
     ):
-        yield api
+        client = mock_client.return_value
+        client.get_devices.return_value = [SnooDevice.from_dict(MOCK_SNOO_DEVICES[0])]
+        client.authorize.return_value = MOCKED_AUTH
+        yield client

tests/components/snoo/const.py

@@ -1,5 +1,9 @@
 """Snoo constants for testing."""
 
+import time
+
+from python_snoo.containers import AuthorizationInfo, SnooData
+
 MOCK_AMAZON_AUTH = {
     # This is a JWT with random values.
     "AccessToken": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJhMWIyYzNkNC1lNWY2"
@@ -32,3 +36,36 @@ MOCK_SNOO_DEVICES = [
         "provisionedAt": "random_time",
     }
 ]
+
+MOCK_SNOO_DATA = SnooData.from_dict(
+    {
+        "system_state": "normal",
+        "sw_version": "v1.14.27",
+        "state_machine": {
+            "session_id": "0",
+            "state": "ONLINE",
+            "is_active_session": "false",
+            "since_session_start_ms": -1,
+            "time_left": -1,
+            "hold": "off",
+            "weaning": "off",
+            "audio": "on",
+            "up_transition": "NONE",
+            "down_transition": "NONE",
+            "sticky_white_noise": "off",
+        },
+        "left_safety_clip": 1,
+        "right_safety_clip": 1,
+        "event": "status_requested",
+        "event_time_ms": int(time.time()),
+        "rx_signal": {"rssi": -45, "strength": 100},
+    }
+)
+
+
+MOCKED_AUTH = AuthorizationInfo(
+    snoo=MOCK_SNOO_AUTH,
+    aws_access=MOCK_AMAZON_AUTH["AccessToken"],
+    aws_id=MOCK_AMAZON_AUTH["IdToken"],
+    aws_refresh=MOCK_AMAZON_AUTH["RefreshToken"],
+)

tests/components/snoo/test_config_flow.py

@@ -13,11 +13,10 @@ from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
 from . import create_entry
-from .conftest import MockedSnoo
 
 
 async def test_config_flow_success(
-    hass: HomeAssistant, mock_setup_entry: AsyncMock, bypass_api: MockedSnoo
+    hass: HomeAssistant, mock_setup_entry: AsyncMock, bypass_api: AsyncMock
 ) -> None:
     """Test we create the entry successfully."""
     result = await hass.config_entries.flow.async_init(
@@ -55,7 +54,7 @@ async def test_config_flow_success(
 async def test_form_auth_issues(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
-    bypass_api: MockedSnoo,
+    bypass_api: AsyncMock,
     exception,
     error_msg,
 ) -> None:
@@ -64,7 +63,7 @@ async def test_form_auth_issues(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     # Set Authorize to fail.
-    bypass_api.set_auth_error(exception)
+    bypass_api.authorize.side_effect = exception
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"],
         {
@@ -73,10 +72,9 @@ async def test_form_auth_issues(
         },
     )
     # Reset auth back to the original
-    bypass_api.set_auth_error(None)
     assert result["type"] == FlowResultType.FORM
     assert result["errors"] == {"base": error_msg}
-
+    bypass_api.authorize.side_effect = None
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"],
         {
@@ -84,7 +82,6 @@ async def test_form_auth_issues(
             CONF_PASSWORD: "test-password",
         },
     )
-    await hass.async_block_till_done()
 
     assert result["type"] == FlowResultType.CREATE_ENTRY
     assert result["title"] == "test-username"
@@ -96,7 +93,7 @@ async def test_form_auth_issues(
 
 
 async def test_account_already_configured(
-    hass: HomeAssistant, mock_setup_entry: AsyncMock, bypass_api
+    hass: HomeAssistant, mock_setup_entry: AsyncMock, bypass_api: AsyncMock
 ) -> None:
     """Ensure we abort if the config flow already exists."""
     create_entry(hass)

tests/components/snoo/test_init.py

@@ -1,14 +1,32 @@
 """Test init for Snoo."""
 
+from unittest.mock import AsyncMock
+
+from python_snoo.exceptions import SnooAuthException
+
+from homeassistant.components.snoo import SnooDeviceError
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
 from . import async_init_integration
-from .conftest import MockedSnoo
 
 
-async def test_async_setup_entry(hass: HomeAssistant, bypass_api: MockedSnoo) -> None:
+async def test_async_setup_entry(hass: HomeAssistant, bypass_api: AsyncMock) -> None:
     """Test a successful setup entry."""
     entry = await async_init_integration(hass)
     assert len(hass.states.async_all("sensor")) == 2
     assert entry.state == ConfigEntryState.LOADED
+
+
+async def test_cannot_auth(hass: HomeAssistant, bypass_api: AsyncMock) -> None:
+    """Test that we are put into retry when we fail to auth."""
+    bypass_api.authorize.side_effect = SnooAuthException
+    entry = await async_init_integration(hass)
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_failed_devices(hass: HomeAssistant, bypass_api: AsyncMock) -> None:
+    """Test that we are put into retry when we fail to get devices."""
+    bypass_api.get_devices.side_effect = SnooDeviceError
+    entry = await async_init_integration(hass)
+    assert entry.state is ConfigEntryState.SETUP_RETRY

tests/components/snoo/test_sensor.py

@@ -0,0 +1,22 @@
+"""Test Snoo Sensors."""
+
+from unittest.mock import AsyncMock
+
+from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+
+from . import async_init_integration, find_update_callback
+from .const import MOCK_SNOO_DATA
+
+
+async def test_sensors(hass: HomeAssistant, bypass_api: AsyncMock) -> None:
+    """Test sensors and check test values are correctly set."""
+    await async_init_integration(hass)
+    assert len(hass.states.async_all("sensor")) == 2
+    assert hass.states.get("sensor.test_snoo_state").state == STATE_UNAVAILABLE
+    assert hass.states.get("sensor.test_snoo_time_left").state == STATE_UNAVAILABLE
+    find_update_callback(bypass_api, "random_num")(MOCK_SNOO_DATA)
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all("sensor")) == 2
+    assert hass.states.get("sensor.test_snoo_state").state == "stop"
+    assert hass.states.get("sensor.test_snoo_time_left").state == STATE_UNKNOWN