Refactor url fetch code to use base platform

homeassistant/components/mqtt/abbreviations.py

@@ -249,6 +249,8 @@ ABBREVIATIONS = {
     "t": "topic",
     "uniq_id": "unique_id",
     "unit_of_meas": "unit_of_measurement",
+    "url_t": "url_topic",
+    "url_tpl": "url_template",
     "val_tpl": "value_template",
     "whit_cmd_t": "white_command_topic",
     "whit_scl": "white_scale",

homeassistant/components/mqtt/image.py

@@ -28,10 +28,10 @@ from homeassistant.util import dt as dt_util
 
 from . import subscription
 from .config import MQTT_BASE_SCHEMA
-from .const import CONF_QOS
+from .const import CONF_ENCODING, CONF_QOS
 from .debug_info import log_messages
 from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
-from .models import ReceiveMessage
+from .models import MessageCallbackType, MqttValueTemplate, ReceiveMessage
 from .util import get_mqtt_data, valid_subscribe_topic
 
 _LOGGER = logging.getLogger(__name__)
@@ -39,21 +39,41 @@ _LOGGER = logging.getLogger(__name__)
 CONF_CONTENT_TYPE = "content_type"
 CONF_IMAGE_ENCODING = "image_encoding"
 CONF_IMAGE_TOPIC = "image_topic"
+CONF_URL_TEMPLATE = "url_template"
+CONF_URL_TOPIC = "url_topic"
 
 DEFAULT_NAME = "MQTT Image"
 
+GET_IMAGE_TIMEOUT = 10
 
-PLATFORM_SCHEMA_MODERN = MQTT_BASE_SCHEMA.extend(
+
+def validate_topic_required(config: ConfigType) -> ConfigType:
+    """Ensure at least one subscribe topic is configured."""
+    if CONF_IMAGE_TOPIC not in config and CONF_URL_TOPIC not in config:
+        raise vol.Invalid("Expected one of [`image_topic`, `url_topic`], got none")
+    if CONF_CONTENT_TYPE in config and CONF_URL_TOPIC in config:
+        raise vol.Invalid(
+            "Option `content_type` can not be used together with `url_topic`"
+        )
+    return config
+
+
+PLATFORM_SCHEMA_BASE = MQTT_BASE_SCHEMA.extend(
     {
-        vol.Optional(CONF_CONTENT_TYPE, default=DEFAULT_CONTENT_TYPE): cv.string,
+        vol.Optional(CONF_CONTENT_TYPE): cv.string,
         vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
-        vol.Required(CONF_IMAGE_TOPIC): valid_subscribe_topic,
+        vol.Exclusive(CONF_URL_TOPIC, "image_topic"): valid_subscribe_topic,
+        vol.Exclusive(CONF_IMAGE_TOPIC, "image_topic"): valid_subscribe_topic,
         vol.Optional(CONF_IMAGE_ENCODING): "b64",
+        vol.Optional(CONF_URL_TEMPLATE): cv.template,
     }
 ).extend(MQTT_ENTITY_COMMON_SCHEMA.schema)
 
+PLATFORM_SCHEMA_MODERN = vol.All(PLATFORM_SCHEMA_BASE.schema, validate_topic_required)
 
-DISCOVERY_SCHEMA = vol.All(PLATFORM_SCHEMA_MODERN.extend({}, extra=vol.REMOVE_EXTRA))
+DISCOVERY_SCHEMA = vol.All(
+    PLATFORM_SCHEMA_BASE.extend({}, extra=vol.REMOVE_EXTRA), validate_topic_required
+)
 
 
 async def async_setup_entry(
@@ -85,6 +105,7 @@ class MqttImage(MqttEntity, ImageEntity):
     _entity_id_format: str = image.ENTITY_ID_FORMAT
     _last_image: bytes | None = None
     _client: httpx.AsyncClient
+    _url: str | None = None
     _url_template: Callable[[ReceivePayloadType], ReceivePayloadType]
     _topic: dict[str, Any]
 
@@ -107,14 +128,43 @@ class MqttImage(MqttEntity, ImageEntity):
 
     def _setup_from_config(self, config: ConfigType) -> None:
         """(Re)Setup the entity."""
-        self._topic = {key: config.get(key) for key in (CONF_IMAGE_TOPIC,)}
-        self._attr_content_type = config[CONF_CONTENT_TYPE]
+        self._topic = {
+            key: config.get(key)
+            for key in (
+                CONF_IMAGE_TOPIC,
+                CONF_URL_TOPIC,
+            )
+        }
+        if CONF_IMAGE_TOPIC in config:
+            self._attr_content_type = config.get(
+                CONF_CONTENT_TYPE, DEFAULT_CONTENT_TYPE
+            )
+        self._url_template = MqttValueTemplate(
+            config.get(CONF_URL_TEMPLATE), entity=self
+        ).async_render_with_possible_json_value
 
     def _prepare_subscribe_topics(self) -> None:
         """(Re)Subscribe to topics."""
 
         topics: dict[str, Any] = {}
 
+        def add_subscribe_topic(topic: str, msg_callback: MessageCallbackType) -> bool:
+            """Add a topic to subscribe to."""
+            encoding: str | None
+            encoding = (
+                None
+                if CONF_IMAGE_TOPIC in self._config
+                else self._config[CONF_ENCODING] or None
+            )
+            if has_topic := self._topic[topic] is not None:
+                topics[topic] = {
+                    "topic": self._topic[topic],
+                    "msg_callback": msg_callback,
+                    "qos": self._config[CONF_QOS],
+                    "encoding": encoding,
+                }
+            return has_topic
+
         @callback
         @log_messages(self.hass, self.entity_id)
         def image_data_received(msg: ReceiveMessage) -> None:
@@ -135,12 +185,27 @@ class MqttImage(MqttEntity, ImageEntity):
             self._attr_image_last_updated = dt_util.utcnow()
             get_mqtt_data(self.hass).state_write_requests.write_state_request(self)
 
-        topics[self._config[CONF_IMAGE_TOPIC]] = {
-            "topic": self._config[CONF_IMAGE_TOPIC],
-            "msg_callback": image_data_received,
-            "qos": self._config[CONF_QOS],
-            "encoding": None,
-        }
+        add_subscribe_topic(CONF_IMAGE_TOPIC, image_data_received)
+
+        @callback
+        @log_messages(self.hass, self.entity_id)
+        def image_from_url_request_received(msg: ReceiveMessage) -> None:
+            """Handle new MQTT messages."""
+
+            try:
+                url = cv.url(self._url_template(msg.payload))
+                self._url = url
+            except vol.Invalid:
+                _LOGGER.error(
+                    "Invalid image URL '%s' received at topic %s",
+                    msg.payload,
+                    msg.topic,
+                )
+                self._last_image = None
+            self._attr_image_last_updated = dt_util.utcnow()
+            get_mqtt_data(self.hass).state_write_requests.write_state_request(self)
+
+        add_subscribe_topic(CONF_URL_TOPIC, image_from_url_request_received)
 
         self._sub_state = subscription.async_prepare_subscribe_topics(
             self.hass, self._sub_state, topics
@@ -152,4 +217,10 @@ class MqttImage(MqttEntity, ImageEntity):
 
     async def async_image(self) -> bytes | None:
         """Return bytes of image."""
-        return self._last_image
+        if CONF_IMAGE_TOPIC in self._config:
+            return self._last_image
+        return await super().async_image()
+
+    async def async_image_url(self) -> str | None:
+        """Return URL of image."""
+        return self._url

tests/components/mqtt/test_image.py

@@ -3,8 +3,10 @@ from base64 import b64encode
 from contextlib import suppress
 from http import HTTPStatus
 import json
-from unittest.mock import patch
+import ssl
+from unittest.mock import MagicMock, patch
 
+import httpx
 import pytest
 import respx
 
@@ -197,11 +199,278 @@ async def test_image_b64_encoded_with_availability(
     assert state.state == "2023-04-01T00:00:00+00:00"
 
 
+@respx.mock
+@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                "image": {
+                    "url_topic": "test/image",
+                    "name": "Test",
+                }
+            }
+        }
+    ],
+)
+async def test_image_from_url(
+    hass: HomeAssistant,
+    hass_client_no_auth: ClientSessionGenerator,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test setup with URL."""
+    respx.get("http://localhost/test.png").respond(
+        status_code=HTTPStatus.OK, content_type="image/png", content=b"milk"
+    )
+    topic = "test/image"
+
+    await mqtt_mock_entry()
+
+    # Test first with invalid URL
+    async_fire_mqtt_message(hass, topic, b"/tmp/test.png")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("image.test")
+    assert state.state == "2023-04-01T00:00:00+00:00"
+
+    assert "Invalid image URL" in caplog.text
+
+    access_token = state.attributes["access_token"]
+    assert state.attributes == {
+        "access_token": access_token,
+        "entity_picture": f"/api/image_proxy/image.test?token={access_token}",
+        "friendly_name": "Test",
+    }
+
+    async_fire_mqtt_message(hass, topic, b"http://localhost/test.png")
+
+    await hass.async_block_till_done()
+
+    client = await hass_client_no_auth()
+    resp = await client.get(state.attributes["entity_picture"])
+    assert resp.status == HTTPStatus.OK
+    body = await resp.text()
+    assert body == "milk"
+
+    state = hass.states.get("image.test")
+    assert state.state == "2023-04-01T00:00:00+00:00"
+
+
+@respx.mock
+@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                "image": {
+                    "url_topic": "test/image",
+                    "name": "Test",
+                    "url_template": "{{ value_json.val }}",
+                }
+            }
+        }
+    ],
+)
+async def test_image_from_url_with_template(
+    hass: HomeAssistant,
+    hass_client_no_auth: ClientSessionGenerator,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test setup with URL."""
+    respx.get("http://localhost/test.png").respond(
+        status_code=HTTPStatus.OK, content_type="image/png", content=b"milk"
+    )
+    topic = "test/image"
+
+    await mqtt_mock_entry()
+
+    state = hass.states.get("image.test")
+    assert state.state == STATE_UNKNOWN
+
+    access_token = state.attributes["access_token"]
+    assert state.attributes == {
+        "access_token": access_token,
+        "entity_picture": f"/api/image_proxy/image.test?token={access_token}",
+        "friendly_name": "Test",
+    }
+
+    async_fire_mqtt_message(hass, topic, '{"val": "http://localhost/test.png"}')
+
+    await hass.async_block_till_done()
+
+    client = await hass_client_no_auth()
+    resp = await client.get(state.attributes["entity_picture"])
+    assert resp.status == HTTPStatus.OK
+    body = await resp.text()
+    assert body == "milk"
+
+    state = hass.states.get("image.test")
+    assert state.state == "2023-04-01T00:00:00+00:00"
+
+
+@respx.mock
+@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                "image": {
+                    "url_topic": "test/image",
+                    "name": "Test",
+                }
+            }
+        }
+    ],
+)
+@pytest.mark.parametrize(
+    ("content_type", "setup_ok"),
+    [
+        ("image/jpg", True),
+        ("image", True),
+        ("image/png", True),
+        ("text/javascript", False),
+    ],
+)
+async def test_image_from_url_content_type(
+    hass: HomeAssistant,
+    hass_client_no_auth: ClientSessionGenerator,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    caplog: pytest.LogCaptureFixture,
+    content_type: str,
+    setup_ok: bool,
+) -> None:
+    """Test setup with URL."""
+    respx.get("http://localhost/test.png").respond(
+        status_code=HTTPStatus.OK, content_type=content_type, content=b"milk"
+    )
+    topic = "test/image"
+
+    await mqtt_mock_entry()
+
+    # Test first with invalid URL
+    async_fire_mqtt_message(hass, topic, b"/tmp/test.png")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("image.test")
+    assert state.state == "2023-04-01T00:00:00+00:00"
+
+    access_token = state.attributes["access_token"]
+    assert state.attributes == {
+        "access_token": access_token,
+        "entity_picture": f"/api/image_proxy/image.test?token={access_token}",
+        "friendly_name": "Test",
+    }
+
+    async_fire_mqtt_message(hass, topic, b"http://localhost/test.png")
+
+    await hass.async_block_till_done()
+
+    client = await hass_client_no_auth()
+    resp = await client.get(state.attributes["entity_picture"])
+    assert resp.status == HTTPStatus.OK if setup_ok else HTTPStatus.SERVICE_UNAVAILABLE
+    if setup_ok:
+        body = await resp.text()
+        assert body == "milk"
+
+    state = hass.states.get("image.test")
+    assert state.state == "2023-04-01T00:00:00+00:00" if setup_ok else STATE_UNKNOWN
+
+
+@respx.mock
+@pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                "image": {
+                    "url_topic": "test/image",
+                    "name": "Test",
+                    "encoding": "utf-8",
+                }
+            }
+        }
+    ],
+)
+@pytest.mark.parametrize(
+    "side_effect",
+    [
+        httpx.RequestError("server offline", request=MagicMock()),
+        httpx.TimeoutException,
+        ssl.SSLError,
+    ],
+)
+async def test_image_from_url_fails(
+    hass: HomeAssistant,
+    hass_client_no_auth: ClientSessionGenerator,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    caplog: pytest.LogCaptureFixture,
+    side_effect: Exception,
+) -> None:
+    """Test setup with minimum configuration."""
+    respx.get("http://localhost/test.png").mock(side_effect=side_effect)
+    topic = "test/image"
+
+    await mqtt_mock_entry()
+
+    state = hass.states.get("image.test")
+    assert state.state == STATE_UNKNOWN
+    access_token = state.attributes["access_token"]
+    assert state.attributes == {
+        "access_token": access_token,
+        "entity_picture": f"/api/image_proxy/image.test?token={access_token}",
+        "friendly_name": "Test",
+    }
+
+    async_fire_mqtt_message(hass, topic, b"http://localhost/test.png")
+
+    await hass.async_block_till_done()
+
+    state = hass.states.get("image.test")
+
+    # The image failed to load, the the last image update is registered
+    # but _last_image was set to `None`
+    assert state.state == "2023-04-01T00:00:00+00:00"
+    client = await hass_client_no_auth()
+    resp = await client.get(state.attributes["entity_picture"])
+    assert resp.status == HTTPStatus.INTERNAL_SERVER_ERROR
+
+
 @respx.mock
 @pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
 @pytest.mark.parametrize(
     ("hass_config", "error_msg"),
     [
+        (
+            {
+                mqtt.DOMAIN: {
+                    "image": {
+                        "url_topic": "test/image",
+                        "content_type": "image/jpg",
+                        "name": "Test",
+                        "encoding": "utf-8",
+                    }
+                }
+            },
+            "Option `content_type` can not be used together with `url_topic`",
+        ),
+        (
+            {
+                mqtt.DOMAIN: {
+                    "image": {
+                        "url_topic": "test/image",
+                        "image_topic": "test/image-data-topic",
+                        "name": "Test",
+                        "encoding": "utf-8",
+                    }
+                }
+            },
+            "two or more values in the same group of exclusion 'image_topic'",
+        ),
         (
             {
                 mqtt.DOMAIN: {
@@ -211,7 +480,7 @@ async def test_image_b64_encoded_with_availability(
                     }
                 }
             },
-            "Invalid config for [mqtt]: required key not provided @ data['mqtt']['image'][0]['image_topic']. Got None.",
+            "Invalid config for [mqtt]: Expected one of [`image_topic`, `url_topic`], got none",
         ),
     ],
 )