Import from top level module in tests

irm_kmi_api/api.py

@@ -14,6 +14,7 @@ from zoneinfo import ZoneInfo
 import aiohttp
 import async_timeout
 
+from .pollen import PollenName, PollenLevel
 from .const import MAP_WARNING_ID_TO_SLUG as SLUG_MAP
 from .const import STYLE_TO_PARAM_MAP, WEEKDAYS, WWEVOL_TO_ENUM_MAP
 from .data import (AnimationFrameData, CurrentWeatherData, ExtendedForecast,
@@ -533,7 +534,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
         return result if len(result) > 0 else []
 
-    async def get_pollen(self) -> Dict[str, str | None]:
+    async def get_pollen(self) -> Dict[PollenName, PollenLevel | None]:
         """
         Get SVG pollen info from the API, return the pollen data dict
 

irm_kmi_api/data.py

@@ -54,8 +54,6 @@ class RadarStyle(Enum):
 
 
 class PollenName(Enum):
-    """Pollens names from the API"""
-
     ALDER = 'alder'
     ASH = 'ash'
     BIRCH = 'birch'

tests/conftest.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 import json
 from unittest.mock import MagicMock
 
-from irm_kmi_api.api import IrmKmiApiClientHa
+from irm_kmi_api import IrmKmiApiClientHa
 from tests.const import IRM_KMI_TO_HA_CONDITION_MAP
 
 

tests/test_api.py

@@ -8,9 +8,9 @@ from zoneinfo import ZoneInfo
 import freezegun
 import pytest
 
-from irm_kmi_api.api import IrmKmiApiClient, IrmKmiApiClientHa
-from irm_kmi_api.data import CurrentWeatherData
-from irm_kmi_api.pollen import PollenParser
+from irm_kmi_api import IrmKmiApiClient, IrmKmiApiClientHa
+from irm_kmi_api import CurrentWeatherData
+from irm_kmi_api import PollenParser
 
 
 @freezegun.freeze_time(dt.fromisoformat('2025-05-03T17:30:00+00:00'))

tests/test_current_weather.py

@@ -4,7 +4,7 @@ from zoneinfo import ZoneInfo
 import pytest
 from freezegun import freeze_time
 
-from irm_kmi_api.data import CurrentWeatherData
+from irm_kmi_api import CurrentWeatherData
 from tests.conftest import get_api_data, get_api_with_data
 from tests.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_PARTLYCLOUDY
 

tests/test_daily_forecast.py

@@ -3,7 +3,7 @@ from zoneinfo import ZoneInfo
 
 from freezegun import freeze_time
 
-from irm_kmi_api.data import ConditionEvol, ExtendedForecast
+from irm_kmi_api import ConditionEvol, ExtendedForecast
 from tests.conftest import get_api_with_data
 from tests.const import ATTR_CONDITION_PARTLYCLOUDY
 

tests/test_hourly_forecast.py

@@ -3,7 +3,7 @@ from zoneinfo import ZoneInfo
 
 from freezegun import freeze_time
 
-from irm_kmi_api.data import Forecast
+from irm_kmi_api import Forecast
 from tests.conftest import get_api_with_data
 from tests.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_RAINY
 

tests/test_pollen.py

@@ -1,8 +1,6 @@
-import logging
 from unittest.mock import AsyncMock
 
-from irm_kmi_api.data import PollenLevel, PollenName
-from irm_kmi_api.pollen import PollenParser
+from irm_kmi_api import PollenLevel, PollenName, PollenParser
 from tests.conftest import get_api_with_data, load_fixture
 
 

tests/test_radar_forecast.py

@@ -1,6 +1,6 @@
 import pytest
 
-from irm_kmi_api.data import RadarForecast
+from irm_kmi_api import RadarForecast
 from tests.conftest import get_api_with_data
 
 

tests/test_rain_graph.py

@@ -6,9 +6,9 @@ from datetime import timedelta
 from unittest.mock import AsyncMock, MagicMock
 from zoneinfo import ZoneInfo
 
-from irm_kmi_api.api import IrmKmiApiClientHa
-from irm_kmi_api.data import AnimationFrameData, RadarAnimationData, RadarStyle
-from irm_kmi_api.rain_graph import RainGraph
+from irm_kmi_api import IrmKmiApiClientHa
+from irm_kmi_api import AnimationFrameData, RadarAnimationData, RadarStyle
+from irm_kmi_api import RainGraph
 from tests.conftest import load_fixture
 
 

tests/test_warning.py

@@ -2,7 +2,7 @@ from datetime import datetime
 
 from freezegun import freeze_time
 
-from irm_kmi_api.data import WarningType
+from irm_kmi_api import WarningType
 from tests.conftest import get_api_with_data