Add Android TV screen capture option and use library screencap (#34074)

homeassistant/components/androidtv/manifest.json

@@ -3,8 +3,8 @@
   "name": "Android TV",
   "documentation": "https://www.home-assistant.io/integrations/androidtv",
   "requirements": [
-    "adb-shell==0.1.1",
-    "androidtv==0.0.39",
+    "adb-shell==0.1.3",
+    "androidtv==0.0.40",
     "pure-python-adb==0.2.2.dev0"
   ],
   "codeowners": ["@JeffLIrion"]

homeassistant/components/androidtv/media_player.py

@@ -1,5 +1,4 @@
 """Support for functionality to interact with Android TV / Fire TV devices."""
-import binascii
 from datetime import datetime
 import functools
 import logging
@@ -89,12 +88,14 @@ CONF_GET_SOURCES = "get_sources"
 CONF_STATE_DETECTION_RULES = "state_detection_rules"
 CONF_TURN_ON_COMMAND = "turn_on_command"
 CONF_TURN_OFF_COMMAND = "turn_off_command"
+CONF_SCREENCAP = "screencap"
 
 DEFAULT_NAME = "Android TV"
 DEFAULT_PORT = 5555
 DEFAULT_ADB_SERVER_PORT = 5037
 DEFAULT_GET_SOURCES = True
 DEFAULT_DEVICE_CLASS = "auto"
+DEFAULT_SCREENCAP = True
 
 DEVICE_ANDROIDTV = "androidtv"
 DEVICE_FIRETV = "firetv"
@@ -146,6 +147,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
             {cv.string: ha_state_detection_rules_validator(vol.Invalid)}
         ),
         vol.Optional(CONF_EXCLUDE_UNNAMED_APPS, default=False): cv.boolean,
+        vol.Optional(CONF_SCREENCAP, default=DEFAULT_SCREENCAP): cv.boolean,
     }
 )
 
@@ -239,6 +241,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
         config.get(CONF_TURN_ON_COMMAND),
         config.get(CONF_TURN_OFF_COMMAND),
         config[CONF_EXCLUDE_UNNAMED_APPS],
+        config[CONF_SCREENCAP],
     ]
 
     if aftv.DEVICE_CLASS == DEVICE_ANDROIDTV:
@@ -382,6 +385,7 @@ class ADBDevice(MediaPlayerDevice):
         turn_on_command,
         turn_off_command,
         exclude_unnamed_apps,
+        screencap,
     ):
         """Initialize the Android TV / Fire TV device."""
         self.aftv = aftv
@@ -401,6 +405,7 @@ class ADBDevice(MediaPlayerDevice):
         self.turn_off_command = turn_off_command
 
         self._exclude_unnamed_apps = exclude_unnamed_apps
+        self._screencap = screencap
 
         # ADB exceptions to catch
         if not self.aftv.adb_server_ip:
@@ -479,7 +484,7 @@ class ADBDevice(MediaPlayerDevice):
 
     async def async_get_media_image(self):
         """Fetch current playing image."""
-        if self.state in [STATE_OFF, None] or not self.available:
+        if not self._screencap or self.state in [STATE_OFF, None] or not self.available:
             return None, None
 
         media_data = await self.hass.async_add_executor_job(self.get_raw_media_data)
@@ -489,16 +494,8 @@ class ADBDevice(MediaPlayerDevice):
 
     @adb_decorator()
     def get_raw_media_data(self):
-        """Raw base64 image data."""
-        try:
-            response = self.aftv.adb_shell("screencap -p | base64")
-        except UnicodeDecodeError:
-            return None
-
-        if isinstance(response, str) and response.strip():
-            return binascii.a2b_base64(response.strip().replace("\n", ""))
-
-        return None
+        """Raw image data."""
+        return self.aftv.adb_screencap()
 
     @property
     def media_image_hash(self):
@@ -613,6 +610,7 @@ class AndroidTVDevice(ADBDevice):
         turn_on_command,
         turn_off_command,
         exclude_unnamed_apps,
+        screencap,
     ):
         """Initialize the Android TV device."""
         super().__init__(
@@ -623,6 +621,7 @@ class AndroidTVDevice(ADBDevice):
             turn_on_command,
             turn_off_command,
             exclude_unnamed_apps,
+            screencap,
         )
 
         self._is_volume_muted = None

requirements_all.txt

@@ -119,7 +119,7 @@ adafruit-circuitpython-bmp280==3.1.1
 adafruit-circuitpython-mcp230xx==2.2.2
 
 # homeassistant.components.androidtv
-adb-shell==0.1.1
+adb-shell==0.1.3
 
 # homeassistant.components.adguard
 adguardhome==0.4.2
@@ -235,7 +235,7 @@ ambiclimate==0.2.1
 amcrest==1.7.0
 
 # homeassistant.components.androidtv
-androidtv==0.0.39
+androidtv==0.0.40
 
 # homeassistant.components.anel_pwrctrl
 anel_pwrctrl-homeassistant==0.0.1.dev2

requirements_test_all.txt

@@ -29,7 +29,7 @@ YesssSMS==0.4.1
 abodepy==0.19.0
 
 # homeassistant.components.androidtv
-adb-shell==0.1.1
+adb-shell==0.1.3
 
 # homeassistant.components.adguard
 adguardhome==0.4.2
@@ -106,7 +106,7 @@ airly==0.0.2
 ambiclimate==0.2.1
 
 # homeassistant.components.androidtv
-androidtv==0.0.39
+androidtv==0.0.40
 
 # homeassistant.components.apns
 apns2==0.3.0

tests/components/androidtv/test_media_player.py

@@ -1,4 +1,5 @@
 """The tests for the androidtv platform."""
+import base64
 import logging
 from unittest.mock import patch
 
@@ -23,6 +24,7 @@ from homeassistant.components.media_player.const import (
     DOMAIN,
     SERVICE_SELECT_SOURCE,
 )
+from homeassistant.components.websocket_api.const import TYPE_RESULT
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     CONF_DEVICE_CLASS,
@@ -968,3 +970,34 @@ async def test_androidtv_volume_set(hass):
         )
 
         patch_set_volume_level.assert_called_with(0.5)
+
+
+async def test_get_image(hass, hass_ws_client):
+    """Test taking a screen capture.
+
+    This is based on `test_get_image` in tests/components/media_player/test_init.py.
+    """
+    patch_key, entity_id = _setup(CONFIG_ANDROIDTV_ADB_SERVER)
+
+    with patchers.PATCH_ADB_DEVICE_TCP, patchers.patch_connect(True)[
+        patch_key
+    ], patchers.patch_shell("")[patch_key]:
+        assert await async_setup_component(hass, DOMAIN, CONFIG_ANDROIDTV_ADB_SERVER)
+
+    with patchers.patch_shell("11")[patch_key]:
+        await hass.helpers.entity_component.async_update_entity(entity_id)
+
+    client = await hass_ws_client(hass)
+
+    with patch("androidtv.basetv.BaseTV.adb_screencap", return_value=b"image"):
+        await client.send_json(
+            {"id": 5, "type": "media_player_thumbnail", "entity_id": entity_id}
+        )
+
+        msg = await client.receive_json()
+
+    assert msg["id"] == 5
+    assert msg["type"] == TYPE_RESULT
+    assert msg["success"]
+    assert msg["result"]["content_type"] == "image/png"
+    assert msg["result"]["content"] == base64.b64encode(b"image").decode("utf-8")