Merge pull request #15 from joostlek/credits

Add credit calculations

Author: joostlek

.github/workflows/linting.yaml

@@ -5,6 +5,7 @@ name: Linting
 on:
   push:
   pull_request:
+    types: [opened, reopened]
   workflow_dispatch:
 
 env:

src/python_opensky/opensky.py

@@ -3,7 +3,8 @@
 
 import asyncio
 import socket
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
 from importlib import metadata
 from typing import Any, cast
 
@@ -24,7 +25,10 @@ class OpenSky:
     session: ClientSession | None = None
     request_timeout: int = 10
     api_host: str = "python_opensky-network.org"
+    opensky_credits: int = 400
+    timezone = timezone.utc
     _close_session: bool = False
+    _credit_usage: dict[datetime, int] = field(default_factory=dict)
 
     async def _request(
         self,
@@ -103,6 +107,7 @@ async def _request(
 
     async def states(self, bounding_box: BoundingBox | None = None) -> StatesResponse:
         """Retrieve state vectors for a given time."""
+        credit_cost = 4
         params = {
             "time": 0,
             "extended": "true",
@@ -114,6 +119,7 @@ async def states(self, bounding_box: BoundingBox | None = None) -> StatesRespons
             params[MAX_LATITUDE] = bounding_box.max_latitude
             params[MIN_LONGITUDE] = bounding_box.min_longitude
             params[MAX_LONGITUDE] = bounding_box.max_longitude
+            credit_cost = self.calculate_credit_costs(bounding_box)
 
         data = await self._request("states/all", data=params)
 
@@ -143,8 +149,37 @@ async def states(self, bounding_box: BoundingBox | None = None) -> StatesRespons
             "states": [dict(zip(keys, state, strict=True)) for state in data["states"]],
         }
 
+        self._register_credit_usage(credit_cost)
+
         return StatesResponse.parse_obj(data)
 
+    @staticmethod
+    def calculate_credit_costs(bounding_box: BoundingBox) -> int:
+        """Calculate the amount of credits a request costs."""
+        latitude_degrees = bounding_box.max_latitude - bounding_box.min_latitude
+        longitude_degrees = bounding_box.max_longitude - bounding_box.min_longitude
+        area = latitude_degrees * longitude_degrees
+        if area < 25:
+            return 1
+        if area < 100:
+            return 2
+        if area < 400:
+            return 3
+        return 4
+
+    def _register_credit_usage(self, opensky_credits: int) -> None:
+        self._credit_usage[datetime.now(self.timezone)] = opensky_credits
+
+    def remaining_credits(self) -> int:
+        """Calculate the remaining opensky credits."""
+        now = datetime.now(self.timezone)
+        used_credits = sum(
+            v
+            for k, v in self._credit_usage.items()
+            if now - timedelta(hours=24) <= k <= now
+        )
+        return self.opensky_credits - used_credits
+
     async def close(self) -> None:
         """Close open client session."""
         if self.session and self._close_session:

tests/test_states.py

@@ -89,6 +89,27 @@ async def test_states_with_bounding_box(
         await opensky.close()
 
 
+async def test_credit_usage(
+    aresponses: ResponsesMockServer,
+) -> None:
+    """Test credit usage."""
+    aresponses.add(
+        OPENSKY_URL,
+        "/api/states/all",
+        "GET",
+        aresponses.Response(
+            status=200,
+            headers={"Content-Type": "application/json"},
+            text=load_fixture("states.json"),
+        ),
+    )
+    async with aiohttp.ClientSession() as session:
+        opensky = OpenSky(session=session)
+        await opensky.states()
+        assert opensky.remaining_credits() == 396
+        await opensky.close()
+
+
 async def test_new_session(
     aresponses: ResponsesMockServer,
 ) -> None:
@@ -173,3 +194,36 @@ async def test_unexpected_server_response(
         with pytest.raises(OpenSkyError):
             assert await opensky.states()
         await opensky.close()
+
+
+async def test_calculating_credit_usage() -> None:
+    """Test calculating credit usage."""
+    opensky = OpenSky()
+    bounding_box = BoundingBox(
+        min_latitude=49.7,
+        max_latitude=50.5,
+        min_longitude=3.2,
+        max_longitude=4.6,
+    )
+    assert opensky.calculate_credit_costs(bounding_box) == 1
+    bounding_box = BoundingBox(
+        min_latitude=46.5,
+        max_latitude=49.9,
+        min_longitude=-1.4,
+        max_longitude=6.8,
+    )
+    assert opensky.calculate_credit_costs(bounding_box) == 2
+    bounding_box = BoundingBox(
+        min_latitude=42.2,
+        max_latitude=49.8,
+        min_longitude=-4.7,
+        max_longitude=10.9,
+    )
+    assert opensky.calculate_credit_costs(bounding_box) == 3
+    bounding_box = BoundingBox(
+        min_latitude=42.2,
+        max_latitude=49.8,
+        min_longitude=-80.7,
+        max_longitude=10.9,
+    )
+    assert opensky.calculate_credit_costs(bounding_box) == 4