Refactor template image (#95353)

pull/95394/head
Erik Montnemery 2023-06-27 18:46:36 +02:00 committed by GitHub
parent 1a8bc1930c
commit 68db751ce1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 126 deletions

View File

@ -4,7 +4,6 @@ from __future__ import annotations
import logging
from typing import Any
import httpx
import voluptuous as vol
from homeassistant.components.image import (
@ -16,7 +15,6 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import TemplateError
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.httpx_client import get_async_client
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
from homeassistant.util import dt as dt_util
@ -82,52 +80,11 @@ async def async_setup_platform(
)
class TemplateImage(ImageEntity):
"""Base class for templated image."""
_last_image: bytes | None = None
_url: str | None = None
_verify_ssl: bool
def __init__(self, hass: HomeAssistant, verify_ssl: bool) -> None:
"""Initialize the image."""
super().__init__(hass)
self._verify_ssl = verify_ssl
async def async_image(self) -> bytes | None:
"""Return bytes of image."""
if self._last_image:
return self._last_image
if not (url := self._url):
return None
try:
async_client = get_async_client(self.hass, verify_ssl=self._verify_ssl)
response = await async_client.get(
url, timeout=GET_IMAGE_TIMEOUT, follow_redirects=True
)
response.raise_for_status()
self._attr_content_type = response.headers["content-type"]
self._last_image = response.content
return self._last_image
except httpx.TimeoutException:
_LOGGER.error("%s: Timeout getting image from %s", self.entity_id, url)
return None
except (httpx.RequestError, httpx.HTTPStatusError) as err:
_LOGGER.error(
"%s: Error getting new image from %s: %s",
self.entity_id,
url,
err,
)
return None
class StateImageEntity(TemplateEntity, TemplateImage):
class StateImageEntity(TemplateEntity, ImageEntity):
"""Representation of a template image."""
_attr_should_poll = False
_attr_image_url: str | None = None
def __init__(
self,
@ -137,7 +94,7 @@ class StateImageEntity(TemplateEntity, TemplateImage):
) -> None:
"""Initialize the image."""
TemplateEntity.__init__(self, hass, config=config, unique_id=unique_id)
TemplateImage.__init__(self, hass, config[CONF_VERIFY_SSL])
ImageEntity.__init__(self, hass, config[CONF_VERIFY_SSL])
self._url_template = config[CONF_URL]
@property
@ -151,11 +108,11 @@ class StateImageEntity(TemplateEntity, TemplateImage):
@callback
def _update_url(self, result):
if isinstance(result, TemplateError):
self._url = None
self._attr_image_url = None
return
self._attr_image_last_updated = dt_util.utcnow()
self._last_image = None
self._url = result
self._cached_image = None
self._attr_image_url = result
async def async_added_to_hass(self) -> None:
"""Register callbacks."""
@ -163,10 +120,10 @@ class StateImageEntity(TemplateEntity, TemplateImage):
await super().async_added_to_hass()
class TriggerImageEntity(TriggerEntity, TemplateImage):
class TriggerImageEntity(TriggerEntity, ImageEntity):
"""Image entity based on trigger data."""
_last_image: bytes | None = None
_attr_image_url: str | None = None
domain = IMAGE_DOMAIN
extra_template_keys = (CONF_URL,)
@ -179,7 +136,7 @@ class TriggerImageEntity(TriggerEntity, TemplateImage):
) -> None:
"""Initialize the entity."""
TriggerEntity.__init__(self, hass, coordinator, config)
TemplateImage.__init__(self, hass, config[CONF_VERIFY_SSL])
ImageEntity.__init__(self, hass, config[CONF_VERIFY_SSL])
@property
def entity_picture(self) -> str | None:
@ -194,5 +151,5 @@ class TriggerImageEntity(TriggerEntity, TemplateImage):
"""Process new data."""
super()._process_data()
self._attr_image_last_updated = dt_util.utcnow()
self._last_image = None
self._url = self._rendered.get(CONF_URL)
self._cached_image = None
self._attr_image_url = self._rendered.get(CONF_URL)

View File

@ -279,78 +279,6 @@ async def test_custom_entity_picture(
)
@respx.mock
@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
async def test_http_error(
hass: HomeAssistant, hass_client: ClientSessionGenerator
) -> None:
"""Test handling http error."""
respx.get("http://example.com").respond(HTTPStatus.NOT_FOUND)
with assert_setup_component(1, "template"):
assert await setup.async_setup_component(
hass,
"template",
{
"template": {
"image": {
"url": "http://example.com",
},
}
},
)
await hass.async_block_till_done()
await hass.async_start()
await hass.async_block_till_done()
expected_state = dt_util.utcnow().isoformat()
await _assert_state(
hass,
hass_client,
expected_state,
b"500: Internal Server Error",
expected_status=HTTPStatus.INTERNAL_SERVER_ERROR,
expected_content_type="text/plain",
)
@respx.mock
@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
async def test_http_timeout(
hass: HomeAssistant, hass_client: ClientSessionGenerator
) -> None:
"""Test handling http timeout."""
respx.get("http://example.com").side_effect = httpx.TimeoutException
with assert_setup_component(1, "template"):
assert await setup.async_setup_component(
hass,
"template",
{
"template": {
"image": {
"url": "http://example.com",
},
}
},
)
await hass.async_block_till_done()
await hass.async_start()
await hass.async_block_till_done()
expected_state = dt_util.utcnow().isoformat()
await _assert_state(
hass,
hass_client,
expected_state,
b"500: Internal Server Error",
expected_status=HTTPStatus.INTERNAL_SERVER_ERROR,
expected_content_type="text/plain",
)
@respx.mock
async def test_template_error(
hass: HomeAssistant, hass_client: ClientSessionGenerator