core/tests/components/roborock/test_image.py

281 lines
9.8 KiB
Python

"""Test Roborock Image platform."""
import copy
from datetime import timedelta
from http import HTTPStatus
import io
from unittest.mock import patch
from PIL import Image
import pytest
from roborock import RoborockException
from vacuum_map_parser_base.map_data import ImageConfig, ImageData
from homeassistant.components.roborock import DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component
from homeassistant.util import dt as dt_util
from .mock_data import MAP_DATA, PROP
from tests.common import MockConfigEntry, async_fire_time_changed
from tests.typing import ClientSessionGenerator
@pytest.fixture
def platforms() -> list[Platform]:
"""Fixture to set platforms used in the test."""
return [Platform.IMAGE]
async def test_floorplan_image(
hass: HomeAssistant,
setup_entry: MockConfigEntry,
hass_client: ClientSessionGenerator,
) -> None:
"""Test floor plan map image is correctly set up."""
assert len(hass.states.async_all("image")) == 4
assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
# Load the image on demand
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
body = await resp.read()
assert body is not None
assert body[0:4] == b"\x89PNG"
# Call a second time - this time forcing it to update - and save new image
now = dt_util.utcnow() + timedelta(minutes=61)
# Copy the device prop so we don't override it
prop = copy.deepcopy(PROP)
prop.status.in_cleaning = 1
new_map_data = copy.deepcopy(MAP_DATA)
new_map_data.image = ImageData(
100, 10, 10, 10, 10, ImageConfig(), Image.new("RGB", (2, 2)), lambda p: p
)
with (
patch(
"homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
return_value=prop,
),
patch(
"homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
),
patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
return_value=new_map_data,
) as parse_map,
):
async_fire_time_changed(hass, now)
await hass.async_block_till_done()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
body = await resp.read()
assert body is not None
assert parse_map.call_count == 1
async def test_floorplan_image_failed_parse(
hass: HomeAssistant,
setup_entry: MockConfigEntry,
hass_client: ClientSessionGenerator,
) -> None:
"""Test that we correctly handle getting None from the image parser."""
client = await hass_client()
map_data = copy.deepcopy(MAP_DATA)
map_data.image = None
now = dt_util.utcnow() + timedelta(seconds=91)
# Copy the device prop so we don't override it
prop = copy.deepcopy(PROP)
prop.status.in_cleaning = 1
# Update image, but get none for parse image.
with (
patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
return_value=map_data,
),
patch(
"homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
return_value=prop,
),
patch(
"homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
),
):
async_fire_time_changed(hass, now)
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert not resp.ok
async def test_load_stored_image(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
setup_entry: MockConfigEntry,
) -> None:
"""Test that we correctly load an image from storage when it already exists."""
img_byte_arr = io.BytesIO()
MAP_DATA.image.data.save(img_byte_arr, format="PNG")
img_bytes = img_byte_arr.getvalue()
# Load the image on demand, which should queue it to be cached on disk
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
with patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
) as parse_map:
# Reload the config entry so that the map is saved in storage and entities exist.
await hass.config_entries.async_reload(setup_entry.entry_id)
await hass.async_block_till_done()
assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
# Test that we can get the image and it correctly serialized and unserialized.
assert resp.status == HTTPStatus.OK
body = await resp.read()
assert body == img_bytes
# Ensure that we never tried to update the map, and only used the cached image.
assert parse_map.call_count == 0
async def test_fail_to_save_image(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
mock_roborock_entry: MockConfigEntry,
bypass_api_fixture,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that we gracefully handle a oserror on saving an image."""
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
# Ensure that map is still working properly.
assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
# Test that we can get the image and it correctly serialized and unserialized.
assert resp.status == HTTPStatus.OK
with patch(
"homeassistant.components.roborock.roborock_storage.Path.write_bytes",
side_effect=OSError,
):
await hass.config_entries.async_unload(mock_roborock_entry.entry_id)
assert "Unable to write map file" in caplog.text
# Config entry is unloaded successfully
assert mock_roborock_entry.state is ConfigEntryState.NOT_LOADED
async def test_fail_to_load_image(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
setup_entry: MockConfigEntry,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that we gracefully handle failing to load an image."""
with (
patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
) as parse_map,
patch(
"homeassistant.components.roborock.roborock_storage.Path.exists",
return_value=True,
),
patch(
"homeassistant.components.roborock.roborock_storage.Path.read_bytes",
side_effect=OSError,
) as read_bytes,
):
# Reload the config entry so that the map is saved in storage and entities exist.
await hass.config_entries.async_reload(setup_entry.entry_id)
await hass.async_block_till_done()
assert read_bytes.call_count == 4
# Ensure that we never updated the map manually since we couldn't load it.
assert parse_map.call_count == 0
assert "Unable to read map file" in caplog.text
async def test_fail_parse_on_startup(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
mock_roborock_entry: MockConfigEntry,
bypass_api_fixture,
) -> None:
"""Test that if we fail parsing on startup, we still create the entity."""
map_data = copy.deepcopy(MAP_DATA)
map_data.image = None
with patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
return_value=map_data,
):
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
assert (
image_entity := hass.states.get("image.roborock_s7_maxv_upstairs")
) is not None
assert image_entity.state
async def test_fail_get_map_on_startup(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
mock_roborock_entry: MockConfigEntry,
bypass_api_fixture,
) -> None:
"""Test that if we fail getting map on startup, we can still create the entity."""
with (
patch(
"homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
return_value=None,
),
):
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
assert (
image_entity := hass.states.get("image.roborock_s7_maxv_upstairs")
) is not None
assert image_entity.state
async def test_fail_updating_image(
hass: HomeAssistant,
setup_entry: MockConfigEntry,
hass_client: ClientSessionGenerator,
) -> None:
"""Test that we handle failing getting the image after it has already been setup.."""
client = await hass_client()
map_data = copy.deepcopy(MAP_DATA)
map_data.image = None
now = dt_util.utcnow() + timedelta(seconds=91)
# Copy the device prop so we don't override it
prop = copy.deepcopy(PROP)
prop.status.in_cleaning = 1
# Update image, but get none for parse image.
with (
patch(
"homeassistant.components.roborock.image.RoborockMapDataParser.parse",
return_value=map_data,
),
patch(
"homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
return_value=prop,
),
patch(
"homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
),
patch(
"homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
side_effect=RoborockException,
),
):
async_fire_time_changed(hass, now)
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert not resp.ok