Add ratelimit state and expires to verify key (#4)

CHANGELOG.md

@@ -1,6 +1,15 @@
 # Changelog
 
-## v0.4.1 (Aug 2023) 
+## v0.4.2 (Aug 2023)
+
+### Additions
+
+- Add `RatelimitState` model.
+- Add `ratelimit` and `expires` fields to `ApiKeyVerification`.
+
+---
+
+## v0.4.1 (Aug 2023)
 
 ### Changes
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unkey.py"
-version = "0.4.1"
+version = "0.4.2"
 description = "An asynchronous Python SDK for unkey.dev."
 authors = ["Jonxslays"]
 license = "GPL-3.0-only"

tests/test_serializer.py

@@ -9,9 +9,6 @@
 from unkey import Serializer
 from unkey import models
 
-# from unittest import mock
-
-
 DictT = t.Dict[str, t.Any]
 
 serializer = Serializer()
@@ -199,6 +196,46 @@ def test_to_api_key(
     assert result == full_api_key
 
 
+####################
+# to_ratelimit_state
+####################
+
+
+def _raw_ratelimit_state() -> DictT:
+    return {
+        "limit": 100,
+        "remaining": 90,
+        "reset": 123,
+    }
+
+
+@pytest.fixture()
+def raw_ratelimit_state() -> DictT:
+    return _raw_ratelimit_state()
+
+
+def _full_ratelimit_state() -> models.RatelimitState:
+    model = models.RatelimitState()
+    model.limit = 100
+    model.remaining = 90
+    model.reset = 123
+    return model
+
+
+@pytest.fixture()
+def full_ratelimit_state() -> models.RatelimitState:
+    return _full_ratelimit_state()
+
+
+def test_to_ratelimit_state(
+    raw_ratelimit_state: DictT,
+    full_ratelimit_state: models.RatelimitState,
+) -> None:
+    result = serializer.to_ratelimit_state(raw_ratelimit_state)
+
+    assert result == full_ratelimit_state
+
+
 #########################
 # to_api_key_verification
 #########################
@@ -210,6 +247,8 @@ def _raw_api_key_verification() -> DictT:
         "owner_id": None,
         "meta": None,
         "remaining": None,
+        "ratelimit": _raw_ratelimit_state(),
+        "expires": 12345,
         "error": "some error",
         "code": "NOT_FOUND",
     }
@@ -226,6 +265,8 @@ def _full_api_key_verification() -> models.ApiKeyVerification:
     model.owner_id = None
     model.meta = None
     model.remaining = None
+    model.ratelimit = _full_ratelimit_state()
+    model.expires = 12345
     model.error = "some error"
     model.code = models.ErrorCode.NotFound
     return model

unkey/__init__.py

@@ -3,7 +3,7 @@
 from typing import Final
 
 __packagename__: Final[str] = "unkey.py"
-__version__: Final[str] = "0.4.1"
+__version__: Final[str] = "0.4.2"
 __author__: Final[str] = "Jonxslays"
 __copyright__: Final[str] = "2023-present Jonxslays"
 __description__: Final[str] = "An asynchronous Python SDK for unkey.dev."
@@ -61,6 +61,7 @@
     "MissingRequiredArgument",
     "Ok",
     "Ratelimit",
+    "RatelimitState",
     "RatelimitType",
     "Result",
     "Route",

unkey/models/__init__.py

@@ -15,5 +15,6 @@
     "ErrorCode",
     "HttpResponse",
     "Ratelimit",
+    "RatelimitState",
     "RatelimitType",
 )

unkey/models/keys.py

@@ -8,7 +8,14 @@
 from .base import BaseModel
 from .http import ErrorCode
 
-__all__ = ("ApiKey", "ApiKeyMeta", "ApiKeyVerification", "Ratelimit", "RatelimitType")
+__all__ = (
+    "ApiKey",
+    "ApiKeyMeta",
+    "ApiKeyVerification",
+    "Ratelimit",
+    "RatelimitState",
+    "RatelimitType",
+)
 
 
 class RatelimitType(BaseEnum):
@@ -107,8 +114,29 @@ class ApiKeyVerification(BaseModel):
     be ignored.
     """
 
+    expires: t.Optional[int]
+    """The unix epoch in milliseconds indicating when this key expires,
+    if it does."""
+
+    ratelimit: t.Optional[RatelimitState]
+    """The state of the ratelimit set on this key, if any."""
+
     code: t.Optional[ErrorCode]
     """The optional error code returned by the unkey api."""
 
     error: t.Optional[str]
     """The error message if the key was invalid."""
+
+
+@attrs.define(init=False, weakref_slot=False)
+class RatelimitState(BaseModel):
+    """The state of the ratelimit for a given key."""
+
+    limit: int
+    """The number of burstable requests allowed."""
+
+    remaining: int
+    """The remaining requests in this burst window."""
+
+    reset: int
+    """The unix timestamp in milliseconds until the next window."""

unkey/serializer.py

@@ -68,13 +68,20 @@ def to_api_key(self, data: DictT) -> models.ApiKey:
 
     def to_api_key_verification(self, data: DictT) -> models.ApiKeyVerification:
         model = models.ApiKeyVerification()
+        ratelimit = data.get("ratelimit")
+        model.ratelimit = self.to_ratelimit_state(ratelimit) if ratelimit else ratelimit
         model.code = models.ErrorCode.from_str_maybe(data.get("code", ""))
         self._set_attrs_cased(
-            model, data, "valid", "owner_id", "meta", "remaining", "error", maybe=True
+            model, data, "valid", "owner_id", "meta", "remaining", "error", "expires", maybe=True
         )
 
         return model
 
+    def to_ratelimit_state(self, data: DictT) -> models.RatelimitState:
+        model = models.RatelimitState()
+        self._set_attrs(model, data, "reset", "limit", "remaining")
+        return model
+
     def to_api(self, data: DictT) -> models.Api:
         model = models.Api()
         self._set_attrs_cased(model, data, "id", "name", "workspace_id")