"""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 HomeAssistant, callback 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: self._async_abort_entries_match( {CONF_HOST: user_input[CONF_HOST], CONF_PORT: user_input[CONF_PORT]} ) # Validate connection to Datadog Agent success = await validate_datadog_connection( self.hass, user_input, ) 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, ) @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: return self.async_show_form( step_id="init", data_schema=vol.Schema( { vol.Required( CONF_PREFIX, default=options.get( CONF_PREFIX, data.get(CONF_PREFIX, DEFAULT_PREFIX) ), ): str, vol.Required( CONF_RATE, default=options.get( CONF_RATE, data.get(CONF_RATE, DEFAULT_RATE) ), ): int, } ), errors={}, ) 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