diff --git a/irm_kmi_api/api.py b/irm_kmi_api/api.py
index 492c594..645f274 100644
--- a/irm_kmi_api/api.py
+++ b/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
 
diff --git a/irm_kmi_api/data.py b/irm_kmi_api/data.py
index 1f8a43e..97f6c95 100644
--- a/irm_kmi_api/data.py
+++ b/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'
diff --git a/tests/conftest.py b/tests/conftest.py
index e207cb0..5d4d5c7 100644
--- a/tests/conftest.py
+++ b/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
 
 
diff --git a/tests/test_api.py b/tests/test_api.py
index 6136d3a..1037f91 100644
--- a/tests/test_api.py
+++ b/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'))
diff --git a/tests/test_current_weather.py b/tests/test_current_weather.py
index adf6ef0..bf07dc6 100644
--- a/tests/test_current_weather.py
+++ b/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
 
diff --git a/tests/test_daily_forecast.py b/tests/test_daily_forecast.py
index c5336bc..5d58560 100644
--- a/tests/test_daily_forecast.py
+++ b/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
 
diff --git a/tests/test_hourly_forecast.py b/tests/test_hourly_forecast.py
index ccee648..9ca0de1 100644
--- a/tests/test_hourly_forecast.py
+++ b/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
 
diff --git a/tests/test_pollen.py b/tests/test_pollen.py
index 86720f3..358b0bc 100644
--- a/tests/test_pollen.py
+++ b/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
 
 
diff --git a/tests/test_radar_forecast.py b/tests/test_radar_forecast.py
index 95b7d0c..a2085c1 100644
--- a/tests/test_radar_forecast.py
+++ b/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
 
 
diff --git a/tests/test_rain_graph.py b/tests/test_rain_graph.py
index c397008..1dbff8a 100644
--- a/tests/test_rain_graph.py
+++ b/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
 
 
diff --git a/tests/test_warning.py b/tests/test_warning.py
index 8f16310..4e0c939 100644
--- a/tests/test_warning.py
+++ b/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