Add Profiler integration (#41175)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

CODEOWNERS

@@ -334,6 +334,7 @@ homeassistant/components/plum_lightpad/* @ColinHarrington @prystupa
 homeassistant/components/point/* @fredrike
 homeassistant/components/poolsense/* @haemishkyd
 homeassistant/components/powerwall/* @bdraco @jrester
+homeassistant/components/profiler/* @bdraco
 homeassistant/components/progettihwsw/* @ardaseremet
 homeassistant/components/prometheus/* @knyar
 homeassistant/components/proxmoxve/* @k4ds3 @jhollowe

homeassistant/components/profiler/__init__.py

@@ -0,0 +1,83 @@
+"""The profiler integration."""
+import asyncio
+import cProfile
+import logging
+import time
+
+from pyprof2calltree import convert
+import voluptuous as vol
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.helpers.service import async_register_admin_service
+from homeassistant.helpers.typing import ConfigType
+
+from .const import DOMAIN
+
+SERVICE_START = "start"
+CONF_SECONDS = "seconds"
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the profiler component."""
+    return True
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
+    """Set up Profiler from a config entry."""
+
+    lock = asyncio.Lock()
+
+    async def _async_run_profile(call: ServiceCall):
+        async with lock:
+            await _async_generate_profile(hass, call)
+
+    async_register_admin_service(
+        hass,
+        DOMAIN,
+        SERVICE_START,
+        _async_run_profile,
+        schema=vol.Schema(
+            {vol.Optional(CONF_SECONDS, default=60.0): vol.Coerce(float)}
+        ),
+    )
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
+    """Unload a config entry."""
+    hass.services.async_remove(domain=DOMAIN, service=SERVICE_START)
+    return True
+
+
+async def _async_generate_profile(hass: HomeAssistant, call: ServiceCall):
+    start_time = int(time.time() * 1000000)
+    hass.components.persistent_notification.async_create(
+        "The profile started. This notification will be updated when it is complete.",
+        title="Profile Started",
+        notification_id=f"profiler_{start_time}",
+    )
+    profiler = cProfile.Profile()
+    profiler.enable()
+    await asyncio.sleep(float(call.data[CONF_SECONDS]))
+    profiler.disable()
+
+    cprofile_path = hass.config.path(f"profile.{start_time}.cprof")
+    callgrind_path = hass.config.path(f"callgrind.out.{start_time}")
+    await hass.async_add_executor_job(
+        _write_profile, profiler, cprofile_path, callgrind_path
+    )
+    hass.components.persistent_notification.async_create(
+        f"Wrote cProfile data to {cprofile_path} and callgrind data to {callgrind_path}",
+        title="Profile Complete",
+        notification_id=f"profiler_{start_time}",
+    )
+
+
+def _write_profile(profiler, cprofile_path, callgrind_path):
+    profiler.create_stats()
+    profiler.dump_stats(cprofile_path)
+    convert(profiler.getstats(), callgrind_path)

homeassistant/components/profiler/config_flow.py

@@ -0,0 +1,28 @@
+"""Config flow for Profiler integration."""
+import logging
+
+import voluptuous as vol
+
+from homeassistant import config_entries
+
+from .const import DEFAULT_NAME
+from .const import DOMAIN  # pylint: disable=unused-import
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for Profiler."""
+
+    VERSION = 1
+    CONNECTION_CLASS = config_entries.CONN_CLASS_UNKNOWN
+
+    async def async_step_user(self, user_input=None):
+        """Handle the initial step."""
+        if self._async_current_entries():
+            return self.async_abort(reason="single_instance_allowed")
+
+        if user_input is not None:
+            return self.async_create_entry(title=DEFAULT_NAME, data={})
+
+        return self.async_show_form(step_id="user", data_schema=vol.Schema({}))

homeassistant/components/profiler/const.py

@@ -0,0 +1,4 @@
+"""Consts used by profiler."""
+
+DOMAIN = "profiler"
+DEFAULT_NAME = "Profiler"

homeassistant/components/profiler/manifest.json

@@ -0,0 +1,13 @@
+{
+  "domain": "profiler",
+  "name": "Profiler",
+  "documentation": "https://www.home-assistant.io/integrations/profiler",
+  "requirements": [
+    "pyprof2calltree==1.4.5"
+  ],
+  "codeowners": [
+    "@bdraco"
+  ],
+  "quality_scale": "internal",
+  "config_flow": true
+}

homeassistant/components/profiler/services.yaml

@@ -0,0 +1,6 @@
+start:
+  description: Start the Profiler
+  fields:
+    seconds:
+      description: The number of seconds to run the profiler.
+      example: 60.0

homeassistant/components/profiler/strings.json

@@ -0,0 +1,12 @@
+{
+  "config": {
+    "step": {
+      "user": {
+        "description": "Are you sure you want to set up the Profiler?"
+      }
+    },
+    "abort": {
+      "single_instance_allowed": "[%key:common::config_flow::abort::single_instance_allowed%]"
+    }
+  }
+}

homeassistant/components/profiler/translations/en.json

@@ -0,0 +1,12 @@
+{
+  "config": {
+    "step": {
+      "user": {
+        "description": "Are you sure you want to set up the Profiler?"
+      }
+    },
+    "abort": {
+      "single_instance_allowed": "[%key:common::config_flow::abort::single_instance_allowed%]"
+    }
+  }
+}

homeassistant/generated/config_flows.py

@@ -145,6 +145,7 @@ FLOWS = [
     "point",
     "poolsense",
     "powerwall",
+    "profiler",
     "progettihwsw",
     "ps4",
     "pvpc_hourly_pricing",

requirements_all.txt

@@ -1588,6 +1588,9 @@ pypjlink2==1.2.1
 # homeassistant.components.point
 pypoint==1.1.2
 
+# homeassistant.components.profiler
+pyprof2calltree==1.4.5
+
 # homeassistant.components.ps4
 pyps4-2ndscreen==1.1.1
 

requirements_test_all.txt

@@ -774,6 +774,9 @@ pyowm==2.10.0
 # homeassistant.components.point
 pypoint==1.1.2
 
+# homeassistant.components.profiler
+pyprof2calltree==1.4.5
+
 # homeassistant.components.ps4
 pyps4-2ndscreen==1.1.1
 

tests/components/profiler/__init__.py

@@ -0,0 +1 @@
+"""Tests for the Profiler integration."""

tests/components/profiler/test_config_flow.py

@@ -0,0 +1,45 @@
+"""Test the Profiler config flow."""
+from homeassistant import config_entries, setup
+from homeassistant.components.profiler.const import DOMAIN
+
+from tests.async_mock import patch
+from tests.common import MockConfigEntry
+
+
+async def test_form_user(hass):
+    """Test we can setup by the user."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "form"
+    assert result["errors"] is None
+
+    with patch(
+        "homeassistant.components.profiler.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.profiler.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+
+    assert result2["type"] == "create_entry"
+    assert result2["title"] == "Profiler"
+    assert result2["data"] == {}
+    await hass.async_block_till_done()
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_form_user_only_once(hass):
+    """Test we can setup by the user only once."""
+    MockConfigEntry(domain=DOMAIN).add_to_hass(hass)
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "abort"
+    assert result["reason"] == "single_instance_allowed"

tests/components/profiler/test_init.py

@@ -0,0 +1,41 @@
+"""Test the Profiler config flow."""
+import os
+
+from homeassistant import setup
+from homeassistant.components.profiler import CONF_SECONDS, SERVICE_START
+from homeassistant.components.profiler.const import DOMAIN
+
+from tests.async_mock import patch
+from tests.common import MockConfigEntry
+
+
+async def test_basic_usage(hass, tmpdir):
+    """Test we can setup and the service is registered."""
+    test_dir = tmpdir.mkdir("profiles")
+
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    entry = MockConfigEntry(domain=DOMAIN)
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert hass.services.has_service(DOMAIN, SERVICE_START)
+
+    last_filename = None
+
+    def _mock_path(filename):
+        nonlocal last_filename
+        last_filename = f"{test_dir}/{filename}"
+        return last_filename
+
+    with patch("homeassistant.components.profiler.cProfile.Profile"), patch.object(
+        hass.config, "path", _mock_path
+    ):
+        await hass.services.async_call(DOMAIN, SERVICE_START, {CONF_SECONDS: 0.000001})
+        await hass.async_block_till_done()
+
+    assert os.path.exists(last_filename)
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()