Fix typing and tests

pyrail/irail.py

@@ -1,12 +1,11 @@
 """Module providing the iRail class for interacting with the iRail API."""
-
 import asyncio
 from asyncio import Lock
 from datetime import datetime
 import logging
 import time
 from types import TracebackType
-from typing import Any, Dict, Type
+from typing import Any, Dict, List, Type
 
 from aiohttp import ClientError, ClientResponse, ClientSession
 
@@ -358,7 +357,7 @@ async def _do_request(self, method: str, args: Dict[str, Any] | None = None) ->
             logger.error("Request failed due to an exception: %s", e)
             return None
 
-    async def get_stations(self) -> list[Station] | None:
+    async def get_stations(self) -> List[Station] | None:
         """Retrieve a list of all train stations from the iRail API.
 
         This method fetches the complete list of available train stations without any additional filtering parameters.
@@ -377,7 +376,8 @@ async def get_stations(self) -> list[Station] | None:
         stations_dict = await self._do_request("stations")
         if stations_dict is None:
             return None
-        stations_response: StationsApiResponse = StationsApiResponse.from_dict(stations_dict)
+        stations_response: StationsApiResponse = StationsApiResponse.from_dict(
+            stations_dict)
         return stations_response.stations
 
     async def get_liveboard(

pyrail/models.py

@@ -1,3 +1,4 @@
+"""Module defining data models for the pyrail application."""
 from dataclasses import dataclass, field
 from typing import List
 
@@ -14,13 +15,18 @@ class ApiResponse(DataClassORJSONMixin):
 @dataclass
 class Station(DataClassORJSONMixin):
     id: str  # The (iRail) ID of the station
-    at_id: str = field(metadata=field_options(alias="@id"))  # Corresponds to "@id" in the schema
-    location_x: float = field(metadata=field_options(alias="locationX"))  # Longitude of the station
-    location_y: float = field(metadata=field_options(alias="locationY"))  # Latitude of the station
-    standard_name: str = field(metadata=field_options(alias="standardname"))  # Consistent name of the station
+    # Corresponds to "@id" in the schema
+    at_id: str = field(metadata=field_options(alias="@id"))
+    location_x: float = field(metadata=field_options(
+        alias="locationX"))  # Longitude of the station
+    location_y: float = field(metadata=field_options(
+        alias="locationY"))  # Latitude of the station
+    standard_name: str = field(metadata=field_options(
+        alias="standardname"))  # Consistent name of the station
     name: str  # Default name of the station
 
 
 @dataclass
 class StationsApiResponse(ApiResponse):
-    stations: List[Station] = field(metadata=field_options(alias="station"))  # List of stations information
+    stations: List[Station] = field(metadata=field_options(
+        alias="station"))  # List of stations information

tests/test_irail.py

@@ -1,12 +1,11 @@
 """Unit tests for the iRail API wrapper."""
-
 from datetime import datetime, timedelta
 from unittest.mock import AsyncMock, patch
 
 from aiohttp import ClientSession
 import pytest
 
-from pyrail.irail import iRail
+from pyrail.irail import Station, iRail
 
 
 @pytest.mark.asyncio
@@ -24,7 +23,8 @@ async def test_successful_request(mock_get):
         mock_get.assert_called_once_with(
             "https://api.irail.be/stations/",
             params={"format": "json", "lang": "en"},
-            headers={"User-Agent": "pyRail (https://github.com/tjorim/pyrail; tielemans.jorim@gmail.com)"},
+            headers={
+                "User-Agent": "pyRail (https://github.com/tjorim/pyrail; tielemans.jorim@gmail.com)"},
         )
         assert response == {"data": "some_data"}
         assert mock_response.status == 200
@@ -54,16 +54,13 @@ async def test_get_stations():
         # Ensure the response is not None
         assert stations is not None, "The response should not be None"
 
-        # Validate that the response is a dictionary
-        assert isinstance(stations, dict), "Expected response to be a dictionary"
-
-        # Validate the presence of key fields
-        assert "station" in stations, "Expected the response to contain a 'station' key"
+        # Validate that the response is a list
+        assert isinstance(stations, list), "Expected response to be a list"
 
-        # Validate the structure of station data
-        station_list = stations.get("station", [])
-        assert isinstance(station_list, list), "Expected 'station' to be a list"
-        assert len(station_list) > 0, "Expected at least one station in the response"
+        # Validate the structure of a station
+        assert isinstance(
+            stations[0], Station), "Expected 'station' to be a list"
+        assert len(stations) > 0, "Expected at least one station in the response"
 
 
 @pytest.mark.asyncio
@@ -83,15 +80,18 @@ async def test_get_connections():
         assert connections is not None, "The response should not be None"
 
         # Validate that the response is a dictionary
-        assert isinstance(connections, dict), "Expected response to be a dictionary"
+        assert isinstance(
+            connections, dict), "Expected response to be a dictionary"
 
         # Validate the presence of key fields
         assert "connection" in connections, "Expected the response to contain a 'connection' key"
 
         # Validate the structure of connection data
         connection_list = connections.get("connection", [])
-        assert isinstance(connection_list, list), "Expected 'connection' to be a list"
-        assert len(connection_list) > 0, "Expected at least one connection in the response"
+        assert isinstance(
+            connection_list, list), "Expected 'connection' to be a list"
+        assert len(
+            connection_list) > 0, "Expected at least one connection in the response"
 
 
 @pytest.mark.asyncio