Fix trafikverket_weatherstation (#60772)

* First commit

* Modify test according to fixes

* Review changes

* Clean up

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

homeassistant/components/trafikverket_weatherstation/__init__.py

@@ -6,19 +6,17 @@ import logging
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN, PLATFORMS
+from .const import PLATFORMS
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Trafikverket Weatherstation from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
-    title = entry.title
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
-    _LOGGER.debug("Loaded entry for %s", title)
+    _LOGGER.debug("Loaded entry for %s", entry.title)
 
     return True
 
@@ -28,12 +26,8 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
-    title = entry.title
     if unload_ok:
-        del hass.data[DOMAIN][entry.entry_id]
+        _LOGGER.debug("Unloaded entry for %s", entry.title)
-        if not hass.data[DOMAIN]:
-            del hass.data[DOMAIN]
-        _LOGGER.debug("Unloaded entry for %s", title)
         return unload_ok
 
     return False

homeassistant/components/trafikverket_weatherstation/config_flow.py

@@ -1,22 +1,18 @@
 """Adds config flow for Trafikverket Weather integration."""
 from __future__ import annotations
 
-import json
+from pytrafikverket.trafikverket_weather import TrafikverketWeather
-
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_API_KEY, CONF_MONITORED_CONDITIONS, CONF_NAME
+from homeassistant.const import CONF_API_KEY
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 
 from .const import CONF_STATION, DOMAIN
-from .sensor import SENSOR_TYPES
-
-SENSOR_LIST: set[str] = {description.key for (description) in SENSOR_TYPES}
 
 DATA_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_NAME): cv.string,
         vol.Required(CONF_API_KEY): cv.string,
         vol.Required(CONF_STATION): cv.string,
     }
@@ -30,14 +26,24 @@ class TVWeatherConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     entry: config_entries.ConfigEntry
 
+    async def validate_input(self, sensor_api: str, station: str) -> str:
+        """Validate input from user input."""
+        web_session = async_get_clientsession(self.hass)
+        weather_api = TrafikverketWeather(web_session, sensor_api)
+        try:
+            await weather_api.async_get_weather(station)
+        except ValueError as err:
+            return str(err)
+        return "connected"
+
     async def async_step_import(self, config: dict):
         """Import a configuration from config.yaml."""
 
         self.context.update(
-            {"title_placeholders": {CONF_NAME: f"YAML import {DOMAIN}"}}
+            {"title_placeholders": {CONF_STATION: f"YAML import {DOMAIN}"}}
         )
 
-        self._async_abort_entries_match({CONF_NAME: config[CONF_NAME]})
+        self._async_abort_entries_match({CONF_STATION: config[CONF_STATION]})
         return await self.async_step_user(user_input=config)
 
     async def async_step_user(self, user_input=None):
@@ -45,20 +51,27 @@ class TVWeatherConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         errors = {}
 
         if user_input is not None:
-            name = user_input[CONF_NAME]
+            name = user_input[CONF_STATION]
             api_key = user_input[CONF_API_KEY]
             station = user_input[CONF_STATION]
-            conditions = json.dumps(list(SENSOR_LIST))
 
-            return self.async_create_entry(
+            validate = await self.validate_input(api_key, station)
-                title=name,
+            if validate == "connected":
-                data={
+                return self.async_create_entry(
-                    CONF_NAME: name,
+                    title=name,
-                    CONF_API_KEY: api_key,
+                    data={
-                    CONF_STATION: station,
+                        CONF_API_KEY: api_key,
-                    CONF_MONITORED_CONDITIONS: conditions,
+                        CONF_STATION: station,
-                },
+                    },
-            )
+                )
+            if validate == "Source: Security, message: Invalid authentication":
+                errors["base"] = "invalid_auth"
+            elif validate == "Could not find a weather station with the specified name":
+                errors["base"] = "invalid_station"
+            elif validate == "Found multiple weather stations with the specified name":
+                errors["base"] = "more_stations"
+            else:
+                errors["base"] = "cannot_connect"
 
         return self.async_show_form(
             step_id="user",

homeassistant/components/trafikverket_weatherstation/sensor.py

@@ -172,7 +172,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Trafikverket sensor entry."""
 
-    sensor_name = entry.data[CONF_NAME]
+    sensor_name = entry.data[CONF_STATION]
     sensor_api = entry.data[CONF_API_KEY]
     sensor_station = entry.data[CONF_STATION]
 
@@ -180,13 +180,11 @@ async def async_setup_entry(
 
     weather_api = TrafikverketWeather(web_session, sensor_api)
 
-    monitored_conditions = entry.data[CONF_MONITORED_CONDITIONS]
     entities = [
         TrafikverketWeatherStation(
             weather_api, sensor_name, sensor_station, description
         )
         for description in SENSOR_TYPES
-        if description.key in monitored_conditions
     ]
 
     async_add_entities(entities, True)

homeassistant/components/trafikverket_weatherstation/strings.json

@@ -3,13 +3,17 @@
       "abort": {
         "already_configured": "[%key:common::config_flow::abort::already_configured_account%]"
       },
+      "error": {
+        "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+        "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
+        "invalid_station": "Could not find a weather station with the specified name",
+        "more_stations": "Found multiple weather stations with the specified name"
+      },
       "step": {
         "user": {
           "data": {
-            "name": "[%key:common::config_flow::data::username%]",
             "api_key": "[%key:common::config_flow::data::api_key%]",
-            "station": "Station",
+            "station": "Station"
-            "conditions": "Monitored conditions"
             }
           }
         }

homeassistant/components/trafikverket_weatherstation/translations/en.json

@@ -3,12 +3,16 @@
         "abort": {
             "already_configured": "Account is already configured"
         },
+        "error": {
+          "cannot_connect": "Failed to connect",
+          "invalid_auth": "Invalid authentication",
+          "invalid_station": "Could not find a weather station with the specified name",
+          "more_stations": "Found multiple weather stations with the specified name"
+        },
         "step": {
             "user": {
                 "data": {
                     "api_key": "API Key",
-                    "conditions": "Monitored conditions",
-                    "name": "Username",
                     "station": "Station"
                 }
             }

tests/components/trafikverket_weatherstation/test_config_flow.py

@@ -1,15 +1,16 @@
 """Test the Trafikverket weatherstation config flow."""
 from __future__ import annotations
 
-import json
 from unittest.mock import patch
 
+import pytest
+
 from homeassistant import config_entries
-from homeassistant.components.trafikverket_weatherstation.sensor import SENSOR_TYPES
 from homeassistant.const import CONF_API_KEY, CONF_NAME
 from homeassistant.core import HomeAssistant
+from homeassistant.data_entry_flow import RESULT_TYPE_FORM
 
-SENSOR_LIST: list[str | None] = {description.key for (description) in SENSOR_TYPES}
+from tests.common import MockConfigEntry
 
 DOMAIN = "trafikverket_weatherstation"
 CONF_STATION = "station"
@@ -25,13 +26,14 @@ async def test_form(hass: HomeAssistant) -> None:
     assert result["errors"] == {}
 
     with patch(
+        "homeassistant.components.trafikverket_weatherstation.config_flow.TrafikverketWeather.async_get_weather",
+    ), patch(
         "homeassistant.components.trafikverket_weatherstation.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
-                CONF_NAME: "Vallby Vasteras",
                 CONF_API_KEY: "1234567890",
                 CONF_STATION: "Vallby",
             },
@@ -39,12 +41,10 @@ async def test_form(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     assert result2["type"] == "create_entry"
-    assert result2["title"] == "Vallby Vasteras"
+    assert result2["title"] == "Vallby"
     assert result2["data"] == {
-        "name": "Vallby Vasteras",
         "api_key": "1234567890",
         "station": "Vallby",
-        "monitored_conditions": json.dumps(list(SENSOR_LIST)),
     }
     assert len(mock_setup_entry.mock_calls) == 1
 
@@ -53,6 +53,8 @@ async def test_import_flow_success(hass: HomeAssistant) -> None:
     """Test a successful import of yaml."""
 
     with patch(
+        "homeassistant.components.trafikverket_weatherstation.config_flow.TrafikverketWeather.async_get_weather",
+    ), patch(
         "homeassistant.components.trafikverket_weatherstation.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
@@ -60,7 +62,7 @@ async def test_import_flow_success(hass: HomeAssistant) -> None:
             DOMAIN,
             context={"source": config_entries.SOURCE_IMPORT},
             data={
-                CONF_NAME: "Vallby Vasteras",
+                CONF_NAME: "Vallby",
                 CONF_API_KEY: "1234567890",
                 CONF_STATION: "Vallby",
             },
@@ -68,11 +70,88 @@ async def test_import_flow_success(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     assert result2["type"] == "create_entry"
-    assert result2["title"] == "Vallby Vasteras"
+    assert result2["title"] == "Vallby"
     assert result2["data"] == {
-        "name": "Vallby Vasteras",
         "api_key": "1234567890",
         "station": "Vallby",
-        "monitored_conditions": json.dumps(list(SENSOR_LIST)),
     }
     assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_import_flow_already_exist(hass: HomeAssistant) -> None:
+    """Test import of yaml already exist."""
+
+    MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_API_KEY: "1234567890",
+            CONF_STATION: "Vallby",
+        },
+    ).add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.trafikverket_weatherstation.async_setup_entry",
+        return_value=True,
+    ), patch(
+        "homeassistant.components.trafikverket_weatherstation.config_flow.TrafikverketWeather.async_get_weather",
+    ):
+        result3 = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_IMPORT},
+            data={
+                CONF_NAME: "Vallby",
+                CONF_API_KEY: "1234567890",
+                CONF_STATION: "Vallby",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result3["type"] == "abort"
+    assert result3["reason"] == "already_configured"
+
+
+@pytest.mark.parametrize(
+    "error_message,base_error",
+    [
+        (
+            "Source: Security, message: Invalid authentication",
+            "invalid_auth",
+        ),
+        (
+            "Could not find a weather station with the specified name",
+            "invalid_station",
+        ),
+        (
+            "Found multiple weather stations with the specified name",
+            "more_stations",
+        ),
+        (
+            "Unknown",
+            "cannot_connect",
+        ),
+    ],
+)
+async def test_flow_fails(
+    hass: HomeAssistant, error_message: str, base_error: str
+) -> None:
+    """Test config flow errors."""
+    result4 = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result4["type"] == RESULT_TYPE_FORM
+    assert result4["step_id"] == config_entries.SOURCE_USER
+
+    with patch(
+        "homeassistant.components.trafikverket_weatherstation.config_flow.TrafikverketWeather.async_get_weather",
+        side_effect=ValueError(error_message),
+    ):
+        result4 = await hass.config_entries.flow.async_configure(
+            result4["flow_id"],
+            user_input={
+                CONF_API_KEY: "1234567890",
+                CONF_STATION: "Vallby",
+            },
+        )
+
+    assert result4["errors"] == {"base": base_error}