From eea22d8079df800581a137512ad0e3b45e86be33 Mon Sep 17 00:00:00 2001 From: Avery <130164016+avedor@users.noreply.github.com> Date: Thu, 24 Jul 2025 06:29:07 -0400 Subject: [PATCH] Add config flow for datadog (#148104) Co-authored-by: G Johansson --- homeassistant/components/datadog/__init__.py | 100 +++++--- .../components/datadog/config_flow.py | 185 ++++++++++++++ homeassistant/components/datadog/const.py | 10 + .../components/datadog/manifest.json | 1 + homeassistant/components/datadog/strings.json | 56 +++++ homeassistant/generated/config_flows.py | 1 + homeassistant/generated/integrations.json | 2 +- tests/components/datadog/common.py | 35 +++ tests/components/datadog/test_config_flow.py | 229 ++++++++++++++++++ tests/components/datadog/test_init.py | 115 +++++++-- 10 files changed, 673 insertions(+), 61 deletions(-) create mode 100644 homeassistant/components/datadog/config_flow.py create mode 100644 homeassistant/components/datadog/const.py create mode 100644 homeassistant/components/datadog/strings.json create mode 100644 tests/components/datadog/common.py create mode 100644 tests/components/datadog/test_config_flow.py diff --git a/homeassistant/components/datadog/__init__.py b/homeassistant/components/datadog/__init__.py index fa852399b09..606f34c9ae0 100644 --- a/homeassistant/components/datadog/__init__.py +++ b/homeassistant/components/datadog/__init__.py @@ -2,9 +2,10 @@ import logging -from datadog import initialize, statsd +from datadog import DogStatsd, initialize import voluptuous as vol +from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry from homeassistant.const import ( CONF_HOST, CONF_PORT, @@ -17,14 +18,19 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv, state as state_helper from homeassistant.helpers.typing import ConfigType +from . import config_flow as config_flow +from .const import ( + CONF_RATE, + DEFAULT_HOST, + DEFAULT_PORT, + DEFAULT_PREFIX, + DEFAULT_RATE, + DOMAIN, +) + _LOGGER = logging.getLogger(__name__) -CONF_RATE = "rate" -DEFAULT_HOST = "localhost" -DEFAULT_PORT = 8125 -DEFAULT_PREFIX = "hass" -DEFAULT_RATE = 1 -DOMAIN = "datadog" +type DatadogConfigEntry = ConfigEntry[DogStatsd] CONFIG_SCHEMA = vol.Schema( { @@ -43,63 +49,85 @@ CONFIG_SCHEMA = vol.Schema( ) -def setup(hass: HomeAssistant, config: ConfigType) -> bool: - """Set up the Datadog component.""" +async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: + """Set up the Datadog integration from YAML, initiating config flow import.""" + if DOMAIN not in config: + return True - conf = config[DOMAIN] - host = conf[CONF_HOST] - port = conf[CONF_PORT] - sample_rate = conf[CONF_RATE] - prefix = conf[CONF_PREFIX] + hass.async_create_task( + hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_IMPORT}, + data=config[DOMAIN], + ) + ) + + return True + + +async def async_setup_entry(hass: HomeAssistant, entry: DatadogConfigEntry) -> bool: + """Set up Datadog from a config entry.""" + + data = entry.data + options = entry.options + host = data[CONF_HOST] + port = data[CONF_PORT] + prefix = options[CONF_PREFIX] + sample_rate = options[CONF_RATE] + + statsd_client = DogStatsd(host=host, port=port, namespace=prefix) + entry.runtime_data = statsd_client initialize(statsd_host=host, statsd_port=port) def logbook_entry_listener(event): - """Listen for logbook entries and send them as events.""" name = event.data.get("name") message = event.data.get("message") - statsd.event( + entry.runtime_data.event( title="Home Assistant", - text=f"%%% \n **{name}** {message} \n %%%", + message=f"%%% \n **{name}** {message} \n %%%", tags=[ f"entity:{event.data.get('entity_id')}", f"domain:{event.data.get('domain')}", ], ) - _LOGGER.debug("Sent event %s", event.data.get("entity_id")) - def state_changed_listener(event): - """Listen for new messages on the bus and sends them to Datadog.""" state = event.data.get("new_state") - if state is None or state.state == STATE_UNKNOWN: return - states = dict(state.attributes) metric = f"{prefix}.{state.domain}" tags = [f"entity:{state.entity_id}"] - for key, value in states.items(): - if isinstance(value, (float, int)): - attribute = f"{metric}.{key.replace(' ', '_')}" + for key, value in state.attributes.items(): + if isinstance(value, (float, int, bool)): value = int(value) if isinstance(value, bool) else value - statsd.gauge(attribute, value, sample_rate=sample_rate, tags=tags) - - _LOGGER.debug("Sent metric %s: %s (tags: %s)", attribute, value, tags) + attribute = f"{metric}.{key.replace(' ', '_')}" + entry.runtime_data.gauge( + attribute, value, sample_rate=sample_rate, tags=tags + ) try: value = state_helper.state_as_number(state) + entry.runtime_data.gauge(metric, value, sample_rate=sample_rate, tags=tags) except ValueError: - _LOGGER.debug("Error sending %s: %s (tags: %s)", metric, state.state, tags) - return + pass - statsd.gauge(metric, value, sample_rate=sample_rate, tags=tags) - - _LOGGER.debug("Sent metric %s: %s (tags: %s)", metric, value, tags) - - hass.bus.listen(EVENT_LOGBOOK_ENTRY, logbook_entry_listener) - hass.bus.listen(EVENT_STATE_CHANGED, state_changed_listener) + entry.async_on_unload( + hass.bus.async_listen(EVENT_LOGBOOK_ENTRY, logbook_entry_listener) + ) + entry.async_on_unload( + hass.bus.async_listen(EVENT_STATE_CHANGED, state_changed_listener) + ) return True + + +async def async_unload_entry(hass: HomeAssistant, entry: DatadogConfigEntry) -> bool: + """Unload a Datadog config entry.""" + runtime = entry.runtime_data + runtime.flush() + runtime.close_socket() + return True diff --git a/homeassistant/components/datadog/config_flow.py b/homeassistant/components/datadog/config_flow.py new file mode 100644 index 00000000000..b4486b0967c --- /dev/null +++ b/homeassistant/components/datadog/config_flow.py @@ -0,0 +1,185 @@ +"""Config flow for Datadog.""" + +from typing import Any + +from datadog import DogStatsd +import voluptuous as vol + +from homeassistant.config_entries import ( + ConfigEntry, + ConfigFlow, + ConfigFlowResult, + OptionsFlow, +) +from homeassistant.const import CONF_HOST, CONF_PORT, CONF_PREFIX +from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant, callback +from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue + +from .const import ( + CONF_RATE, + DEFAULT_HOST, + DEFAULT_PORT, + DEFAULT_PREFIX, + DEFAULT_RATE, + DOMAIN, +) + + +class DatadogConfigFlow(ConfigFlow, domain=DOMAIN): + """Handle a config flow for Datadog.""" + + VERSION = 1 + + async def async_step_user( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Handle user config flow.""" + errors: dict[str, str] = {} + if user_input: + # Validate connection to Datadog Agent + success = await validate_datadog_connection( + self.hass, + user_input, + ) + self._async_abort_entries_match( + {CONF_HOST: user_input[CONF_HOST], CONF_PORT: user_input[CONF_PORT]} + ) + if not success: + errors["base"] = "cannot_connect" + else: + return self.async_create_entry( + title=f"Datadog {user_input['host']}", + data={ + CONF_HOST: user_input[CONF_HOST], + CONF_PORT: user_input[CONF_PORT], + }, + options={ + CONF_PREFIX: user_input[CONF_PREFIX], + CONF_RATE: user_input[CONF_RATE], + }, + ) + + return self.async_show_form( + step_id="user", + data_schema=vol.Schema( + { + vol.Required(CONF_HOST, default=DEFAULT_HOST): str, + vol.Required(CONF_PORT, default=DEFAULT_PORT): int, + vol.Required(CONF_PREFIX, default=DEFAULT_PREFIX): str, + vol.Required(CONF_RATE, default=DEFAULT_RATE): int, + } + ), + errors=errors, + ) + + async def async_step_import(self, user_input: dict[str, Any]) -> ConfigFlowResult: + """Handle import from configuration.yaml.""" + # Check for duplicates + self._async_abort_entries_match( + {CONF_HOST: user_input[CONF_HOST], CONF_PORT: user_input[CONF_PORT]} + ) + + result = await self.async_step_user(user_input) + + if errors := result.get("errors"): + await deprecate_yaml_issue(self.hass, False) + return self.async_abort(reason=errors["base"]) + + await deprecate_yaml_issue(self.hass, True) + return result + + @staticmethod + @callback + def async_get_options_flow(config_entry: ConfigEntry) -> OptionsFlow: + """Get the options flow handler.""" + return DatadogOptionsFlowHandler() + + +class DatadogOptionsFlowHandler(OptionsFlow): + """Handle Datadog options.""" + + async def async_step_init( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Manage the Datadog options.""" + errors: dict[str, str] = {} + data = self.config_entry.data + options = self.config_entry.options + + if user_input is None: + user_input = {} + + success = await validate_datadog_connection( + self.hass, + {**data, **user_input}, + ) + if success: + return self.async_create_entry( + data={ + CONF_PREFIX: user_input[CONF_PREFIX], + CONF_RATE: user_input[CONF_RATE], + } + ) + + errors["base"] = "cannot_connect" + return self.async_show_form( + step_id="init", + data_schema=vol.Schema( + { + vol.Required(CONF_PREFIX, default=options[CONF_PREFIX]): str, + vol.Required(CONF_RATE, default=options[CONF_RATE]): int, + } + ), + errors=errors, + ) + + +async def validate_datadog_connection( + hass: HomeAssistant, user_input: dict[str, Any] +) -> bool: + """Attempt to send a test metric to the Datadog agent.""" + try: + client = DogStatsd(user_input[CONF_HOST], user_input[CONF_PORT]) + await hass.async_add_executor_job(client.increment, "connection_test") + except (OSError, ValueError): + return False + else: + return True + + +async def deprecate_yaml_issue( + hass: HomeAssistant, + import_success: bool, +) -> None: + """Create an issue to deprecate YAML config.""" + if import_success: + async_create_issue( + hass, + HOMEASSISTANT_DOMAIN, + f"deprecated_yaml_{DOMAIN}", + is_fixable=False, + issue_domain=DOMAIN, + breaks_in_ha_version="2026.2.0", + severity=IssueSeverity.WARNING, + translation_key="deprecated_yaml", + translation_placeholders={ + "domain": DOMAIN, + "integration_title": "Datadog", + }, + ) + else: + async_create_issue( + hass, + DOMAIN, + "deprecated_yaml_import_connection_error", + breaks_in_ha_version="2026.2.0", + is_fixable=False, + issue_domain=DOMAIN, + severity=IssueSeverity.WARNING, + translation_key="deprecated_yaml_import_connection_error", + translation_placeholders={ + "domain": DOMAIN, + "integration_title": "Datadog", + "url": f"/config/integrations/dashboard/add?domain={DOMAIN}", + }, + ) diff --git a/homeassistant/components/datadog/const.py b/homeassistant/components/datadog/const.py new file mode 100644 index 00000000000..e9e5d80eeba --- /dev/null +++ b/homeassistant/components/datadog/const.py @@ -0,0 +1,10 @@ +"""Constants for the Datadog integration.""" + +DOMAIN = "datadog" + +CONF_RATE = "rate" + +DEFAULT_HOST = "localhost" +DEFAULT_PORT = 8125 +DEFAULT_PREFIX = "hass" +DEFAULT_RATE = 1 diff --git a/homeassistant/components/datadog/manifest.json b/homeassistant/components/datadog/manifest.json index ca9681effca..815446b9ab4 100644 --- a/homeassistant/components/datadog/manifest.json +++ b/homeassistant/components/datadog/manifest.json @@ -2,6 +2,7 @@ "domain": "datadog", "name": "Datadog", "codeowners": [], + "config_flow": true, "documentation": "https://www.home-assistant.io/integrations/datadog", "iot_class": "local_push", "loggers": ["datadog"], diff --git a/homeassistant/components/datadog/strings.json b/homeassistant/components/datadog/strings.json new file mode 100644 index 00000000000..86bb2019fc1 --- /dev/null +++ b/homeassistant/components/datadog/strings.json @@ -0,0 +1,56 @@ +{ + "config": { + "step": { + "user": { + "description": "Enter your Datadog Agent's address and port.", + "data": { + "host": "[%key:common::config_flow::data::host%]", + "port": "[%key:common::config_flow::data::port%]", + "prefix": "Prefix", + "rate": "Rate" + }, + "data_description": { + "host": "The hostname or IP address of the Datadog Agent.", + "port": "Port the Datadog Agent is listening on", + "prefix": "Metric prefix to use", + "rate": "The sample rate of UDP packets sent to Datadog." + } + } + }, + "error": { + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]" + }, + "abort": { + "already_configured": "[%key:common::config_flow::abort::already_configured_service%]", + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]" + } + }, + "options": { + "step": { + "init": { + "description": "Update the Datadog configuration.", + "data": { + "prefix": "[%key:component::datadog::config::step::user::data::prefix%]", + "rate": "[%key:component::datadog::config::step::user::data::rate%]" + }, + "data_description": { + "prefix": "[%key:component::datadog::config::step::user::data_description::prefix%]", + "rate": "[%key:component::datadog::config::step::user::data_description::rate%]" + } + } + }, + "error": { + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]" + }, + "abort": { + "already_configured": "[%key:common::config_flow::abort::already_configured_service%]", + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]" + } + }, + "issues": { + "deprecated_yaml_import_connection_error": { + "title": "{domain} YAML configuration import failed", + "description": "There was an error connecting to the Datadog Agent when trying to import the YAML configuration.\n\nEnsure the YAML configuration is correct and restart Home Assistant to try again or remove the {domain} configuration from your `configuration.yaml` file and continue to [set up the integration]({url}) manually." + } + } +} diff --git a/homeassistant/generated/config_flows.py b/homeassistant/generated/config_flows.py index 49695b695ac..d9fd32d204b 100644 --- a/homeassistant/generated/config_flows.py +++ b/homeassistant/generated/config_flows.py @@ -124,6 +124,7 @@ FLOWS = { "cpuspeed", "crownstone", "daikin", + "datadog", "deako", "deconz", "deluge", diff --git a/homeassistant/generated/integrations.json b/homeassistant/generated/integrations.json index 431ece3f81a..33cc637b8a8 100644 --- a/homeassistant/generated/integrations.json +++ b/homeassistant/generated/integrations.json @@ -1171,7 +1171,7 @@ "datadog": { "name": "Datadog", "integration_type": "hub", - "config_flow": false, + "config_flow": true, "iot_class": "local_push" }, "ddwrt": { diff --git a/tests/components/datadog/common.py b/tests/components/datadog/common.py new file mode 100644 index 00000000000..07539dc0e07 --- /dev/null +++ b/tests/components/datadog/common.py @@ -0,0 +1,35 @@ +"""Common helpers for the datetime entity component tests.""" + +from unittest import mock + +MOCK_DATA = { + "host": "localhost", + "port": 8125, +} + +MOCK_OPTIONS = { + "prefix": "hass", + "rate": 1, +} + +MOCK_CONFIG = {**MOCK_DATA, **MOCK_OPTIONS} + +MOCK_YAML_INVALID = { + "host": "127.0.0.1", + "port": 65535, + "prefix": "failtest", + "rate": 1, +} + + +CONNECTION_TEST_METRIC = "connection_test" + + +def create_mock_state(entity_id, state, attributes=None): + """Helper to create a mock state object.""" + mock_state = mock.MagicMock() + mock_state.entity_id = entity_id + mock_state.state = state + mock_state.domain = entity_id.split(".")[0] + mock_state.attributes = attributes or {} + return mock_state diff --git a/tests/components/datadog/test_config_flow.py b/tests/components/datadog/test_config_flow.py new file mode 100644 index 00000000000..7950bb2c17d --- /dev/null +++ b/tests/components/datadog/test_config_flow.py @@ -0,0 +1,229 @@ +"""Tests for the Datadog config flow.""" + +from unittest.mock import MagicMock, patch + +from homeassistant.components import datadog +from homeassistant.config_entries import SOURCE_IMPORT +from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant +from homeassistant.data_entry_flow import FlowResultType +import homeassistant.helpers.issue_registry as ir + +from .common import MOCK_CONFIG, MOCK_DATA, MOCK_OPTIONS, MOCK_YAML_INVALID + +from tests.common import MockConfigEntry + + +async def test_user_flow_success(hass: HomeAssistant) -> None: + """Test user-initiated config flow.""" + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd: + mock_instance = MagicMock() + mock_dogstatsd.return_value = mock_instance + + result = await hass.config_entries.flow.async_init( + datadog.DOMAIN, context={"source": "user"} + ) + assert result["type"] == FlowResultType.FORM + + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], user_input=MOCK_CONFIG + ) + assert result2["title"] == f"Datadog {MOCK_CONFIG['host']}" + assert result2["type"] == FlowResultType.CREATE_ENTRY + assert result2["data"] == MOCK_DATA + assert result2["options"] == MOCK_OPTIONS + + +async def test_user_flow_retry_after_connection_fail(hass: HomeAssistant) -> None: + """Test connection failure.""" + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + side_effect=OSError("Connection failed"), + ): + result = await hass.config_entries.flow.async_init( + datadog.DOMAIN, context={"source": "user"} + ) + + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], user_input=MOCK_CONFIG + ) + assert result2["type"] == FlowResultType.FORM + assert result2["errors"] == {"base": "cannot_connect"} + + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + ): + result3 = await hass.config_entries.flow.async_configure( + result["flow_id"], user_input=MOCK_CONFIG + ) + assert result3["type"] == FlowResultType.CREATE_ENTRY + assert result3["data"] == MOCK_DATA + assert result3["options"] == MOCK_OPTIONS + + +async def test_options_flow_cannot_connect(hass: HomeAssistant) -> None: + """Test that the options flow shows an error when connection fails.""" + mock_entry = MockConfigEntry( + domain=datadog.DOMAIN, + data=MOCK_DATA, + options=MOCK_OPTIONS, + ) + mock_entry.add_to_hass(hass) + + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + side_effect=OSError("connection failed"), + ): + result = await hass.config_entries.options.async_init(mock_entry.entry_id) + assert result["type"] == FlowResultType.FORM + + result2 = await hass.config_entries.options.async_configure( + result["flow_id"], user_input=MOCK_OPTIONS + ) + + assert result2["type"] == FlowResultType.FORM + assert result2["errors"] == {"base": "cannot_connect"} + + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + ): + result3 = await hass.config_entries.options.async_configure( + result["flow_id"], user_input=MOCK_OPTIONS + ) + assert result3["type"] == FlowResultType.CREATE_ENTRY + assert result3["data"] == MOCK_OPTIONS + + +async def test_import_flow( + hass: HomeAssistant, issue_registry: ir.IssueRegistry +) -> None: + """Test import triggers config flow and is accepted.""" + with ( + patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd, + ): + mock_instance = MagicMock() + mock_dogstatsd.return_value = mock_instance + + result = await hass.config_entries.flow.async_init( + datadog.DOMAIN, + context={"source": SOURCE_IMPORT}, + data=MOCK_CONFIG, + ) + + assert result["type"] == FlowResultType.CREATE_ENTRY + assert result["data"] == MOCK_DATA + assert result["options"] == MOCK_OPTIONS + + await hass.async_block_till_done() + + # Deprecation issue should be created + issue = issue_registry.async_get_issue( + HOMEASSISTANT_DOMAIN, "deprecated_yaml_datadog" + ) + assert issue is not None + assert issue.translation_key == "deprecated_yaml" + assert issue.severity == ir.IssueSeverity.WARNING + + +async def test_import_connection_error( + hass: HomeAssistant, issue_registry: ir.IssueRegistry +) -> None: + """Test import triggers connection error issue.""" + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + side_effect=OSError("connection refused"), + ): + result = await hass.config_entries.flow.async_init( + datadog.DOMAIN, + context={"source": SOURCE_IMPORT}, + data=MOCK_YAML_INVALID, + ) + assert result["type"] == "abort" + assert result["reason"] == "cannot_connect" + + issue = issue_registry.async_get_issue( + datadog.DOMAIN, "deprecated_yaml_import_connection_error" + ) + assert issue is not None + assert issue.translation_key == "deprecated_yaml_import_connection_error" + assert issue.severity == ir.IssueSeverity.WARNING + + +async def test_options_flow(hass: HomeAssistant) -> None: + """Test updating options after setup.""" + mock_entry = MockConfigEntry( + domain=datadog.DOMAIN, + data=MOCK_DATA, + options=MOCK_OPTIONS, + ) + mock_entry.add_to_hass(hass) + + new_options = { + "prefix": "updated", + "rate": 5, + } + + # OSError Case + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + side_effect=OSError, + ): + result = await hass.config_entries.options.async_init(mock_entry.entry_id) + assert result["type"] == FlowResultType.FORM + result2 = await hass.config_entries.options.async_configure( + result["flow_id"], user_input=new_options + ) + assert result2["type"] == FlowResultType.FORM + assert result2["errors"] == {"base": "cannot_connect"} + + # ValueError Case + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd", + side_effect=ValueError, + ): + result = await hass.config_entries.options.async_init(mock_entry.entry_id) + assert result["type"] == FlowResultType.FORM + result2 = await hass.config_entries.options.async_configure( + result["flow_id"], user_input=new_options + ) + assert result2["type"] == FlowResultType.FORM + assert result2["errors"] == {"base": "cannot_connect"} + + # Success Case + with patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd: + mock_instance = MagicMock() + mock_dogstatsd.return_value = mock_instance + + result = await hass.config_entries.options.async_configure( + result["flow_id"], user_input=new_options + ) + + assert result["type"] == FlowResultType.CREATE_ENTRY + assert result["data"] == new_options + mock_instance.increment.assert_called_once_with("connection_test") + + +async def test_import_flow_abort_already_configured_service( + hass: HomeAssistant, +) -> None: + """Abort import if the same host/port is already configured.""" + existing_entry = MockConfigEntry( + domain=datadog.DOMAIN, + data=MOCK_DATA, + options=MOCK_OPTIONS, + ) + existing_entry.add_to_hass(hass) + + result = await hass.config_entries.flow.async_init( + datadog.DOMAIN, + context={"source": "import"}, + data=MOCK_CONFIG, + ) + + assert result["type"] == "abort" + assert result["reason"] == "already_configured" diff --git a/tests/components/datadog/test_init.py b/tests/components/datadog/test_init.py index 3b7bea3c926..73bce96d16c 100644 --- a/tests/components/datadog/test_init.py +++ b/tests/components/datadog/test_init.py @@ -4,11 +4,15 @@ from unittest import mock from unittest.mock import patch from homeassistant.components import datadog -from homeassistant.const import EVENT_LOGBOOK_ENTRY, STATE_OFF, STATE_ON +from homeassistant.components.datadog import async_setup_entry +from homeassistant.config_entries import ConfigEntryState +from homeassistant.const import EVENT_LOGBOOK_ENTRY, STATE_OFF, STATE_ON, STATE_UNKNOWN from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component -from tests.common import assert_setup_component +from .common import MOCK_DATA, MOCK_OPTIONS, create_mock_state + +from tests.common import EVENT_STATE_CHANGED, MockConfigEntry, assert_setup_component async def test_invalid_config(hass: HomeAssistant) -> None: @@ -24,20 +28,22 @@ async def test_datadog_setup_full(hass: HomeAssistant) -> None: config = {datadog.DOMAIN: {"host": "host", "port": 123, "rate": 1, "prefix": "foo"}} with ( - patch("homeassistant.components.datadog.initialize") as mock_init, - patch("homeassistant.components.datadog.statsd"), + patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd, ): assert await async_setup_component(hass, datadog.DOMAIN, config) - assert mock_init.call_count == 1 - assert mock_init.call_args == mock.call(statsd_host="host", statsd_port=123) + assert mock_dogstatsd.call_count == 1 + assert mock_dogstatsd.call_args == mock.call("host", 123) async def test_datadog_setup_defaults(hass: HomeAssistant) -> None: """Test setup with defaults.""" with ( - patch("homeassistant.components.datadog.initialize") as mock_init, - patch("homeassistant.components.datadog.statsd"), + patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd, ): assert await async_setup_component( hass, @@ -51,20 +57,31 @@ async def test_datadog_setup_defaults(hass: HomeAssistant) -> None: }, ) - assert mock_init.call_count == 1 - assert mock_init.call_args == mock.call(statsd_host="host", statsd_port=8125) + assert mock_dogstatsd.call_count == 1 + assert mock_dogstatsd.call_args == mock.call("host", 8125) async def test_logbook_entry(hass: HomeAssistant) -> None: """Test event listener.""" with ( - patch("homeassistant.components.datadog.initialize"), - patch("homeassistant.components.datadog.statsd") as mock_statsd, + patch("homeassistant.components.datadog.DogStatsd") as mock_statsd_class, + patch( + "homeassistant.components.datadog.config_flow.DogStatsd", mock_statsd_class + ), ): + mock_statsd = mock_statsd_class.return_value + assert await async_setup_component( hass, datadog.DOMAIN, - {datadog.DOMAIN: {"host": "host", "rate": datadog.DEFAULT_RATE}}, + { + datadog.DOMAIN: { + "host": "host", + "port": datadog.DEFAULT_PORT, + "rate": datadog.DEFAULT_RATE, + "prefix": datadog.DEFAULT_PREFIX, + } + }, ) event = { @@ -79,19 +96,21 @@ async def test_logbook_entry(hass: HomeAssistant) -> None: assert mock_statsd.event.call_count == 1 assert mock_statsd.event.call_args == mock.call( title="Home Assistant", - text=f"%%% \n **{event['name']}** {event['message']} \n %%%", + message=f"%%% \n **{event['name']}** {event['message']} \n %%%", tags=["entity:sensor.foo.bar", "domain:automation"], ) - mock_statsd.event.reset_mock() - async def test_state_changed(hass: HomeAssistant) -> None: """Test event listener.""" with ( - patch("homeassistant.components.datadog.initialize"), - patch("homeassistant.components.datadog.statsd") as mock_statsd, + patch("homeassistant.components.datadog.DogStatsd") as mock_statsd_class, + patch( + "homeassistant.components.datadog.config_flow.DogStatsd", mock_statsd_class + ), ): + mock_statsd = mock_statsd_class.return_value + assert await async_setup_component( hass, datadog.DOMAIN, @@ -109,12 +128,7 @@ async def test_state_changed(hass: HomeAssistant) -> None: attributes = {"elevation": 3.2, "temperature": 5.0, "up": True, "down": False} for in_, out in valid.items(): - state = mock.MagicMock( - domain="sensor", - entity_id="sensor.foobar", - state=in_, - attributes=attributes, - ) + state = create_mock_state("sensor.foobar", in_, attributes) hass.states.async_set(state.entity_id, state.state, state.attributes) await hass.async_block_till_done() assert mock_statsd.gauge.call_count == 5 @@ -145,3 +159,56 @@ async def test_state_changed(hass: HomeAssistant) -> None: hass.states.async_set("domain.test", invalid, {}) await hass.async_block_till_done() assert not mock_statsd.gauge.called + + +async def test_unload_entry(hass: HomeAssistant) -> None: + """Test unloading the config entry cleans up properly.""" + client = mock.MagicMock() + + with ( + patch("homeassistant.components.datadog.DogStatsd", return_value=client), + patch("homeassistant.components.datadog.initialize"), + ): + entry = MockConfigEntry( + domain=datadog.DOMAIN, + data=MOCK_DATA, + options=MOCK_OPTIONS, + ) + entry.add_to_hass(hass) + + assert await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + + assert entry.state is ConfigEntryState.LOADED + assert await hass.config_entries.async_unload(entry.entry_id) + await hass.async_block_till_done() + + assert entry.state is ConfigEntryState.NOT_LOADED + + client.flush.assert_called_once() + client.close_socket.assert_called_once() + + +async def test_state_changed_skips_unknown(hass: HomeAssistant) -> None: + """Test state_changed_listener skips None and unknown states.""" + entry = MockConfigEntry(domain=datadog.DOMAIN, data=MOCK_DATA, options=MOCK_OPTIONS) + entry.add_to_hass(hass) + + with ( + patch( + "homeassistant.components.datadog.config_flow.DogStatsd" + ) as mock_dogstatsd, + ): + await async_setup_entry(hass, entry) + + # Test None state + hass.bus.async_fire(EVENT_STATE_CHANGED, {"new_state": None}) + await hass.async_block_till_done() + assert not mock_dogstatsd.gauge.called + + # Test STATE_UNKNOWN + unknown_state = mock.MagicMock() + unknown_state.state = STATE_UNKNOWN + hass.bus.async_fire(EVENT_STATE_CHANGED, {"new_state": unknown_state}) + await hass.async_block_till_done() + assert not mock_dogstatsd.gauge.called