Merge pull request #89 from jcxldn/feat/jcx/rate-limits

Add support for microsoft-defined rate limits

Author: tuxuser

tests/test_ratelimits.py

@@ -0,0 +1,241 @@
+from datetime import datetime, timedelta
+from httpx import Response
+import pytest
+import asyncio
+
+from tests.common import get_response_json
+from xbox.webapi.api.provider.ratelimitedprovider import RateLimitedProvider
+
+from xbox.webapi.common.exceptions import RateLimitExceededException, XboxException
+from xbox.webapi.common.ratelimits import CombinedRateLimit
+from xbox.webapi.common.ratelimits.models import TimePeriod
+
+
+def helper_test_combinedratelimit(
+    crl: CombinedRateLimit, burstLimit: int, sustainLimit: int
+):
+    burst = crl.get_limits_by_period(TimePeriod.BURST)
+    sustain = crl.get_limits_by_period(TimePeriod.SUSTAIN)
+
+    # These functions should return a list with one element
+    assert type(burst) == list
+    assert type(sustain) == list
+
+    assert len(burst) == 1
+    assert len(sustain) == 1
+
+    # Check that their limits are what we expect
+    assert burst[0].get_limit() == burstLimit
+    assert sustain[0].get_limit() == sustainLimit
+
+
+def test_ratelimitedprovider_rate_limits_same_rw_values(xbl_client):
+    class child_class(RateLimitedProvider):
+        RATE_LIMITS = {"burst": 1, "sustain": 2}
+
+    instance = child_class(xbl_client)
+
+    helper_test_combinedratelimit(instance.rate_limit_read, 1, 2)
+    helper_test_combinedratelimit(instance.rate_limit_write, 1, 2)
+
+
+def test_ratelimitedprovider_rate_limits_diff_rw_values(xbl_client):
+    class child_class(RateLimitedProvider):
+        RATE_LIMITS = {
+            "burst": {"read": 1, "write": 2},
+            "sustain": {"read": 3, "write": 4},
+        }
+
+    instance = child_class(xbl_client)
+
+    helper_test_combinedratelimit(instance.rate_limit_read, 1, 3)
+    helper_test_combinedratelimit(instance.rate_limit_write, 2, 4)
+
+
+def test_ratelimitedprovider_rate_limits_mixed(xbl_client):
+    class burst_diff(RateLimitedProvider):
+        RATE_LIMITS = {"burst": {"read": 1, "write": 2}, "sustain": 3}
+
+    burst_diff_inst = burst_diff(xbl_client)
+
+    # Sustain values are the same (third paramater)
+    helper_test_combinedratelimit(burst_diff_inst.rate_limit_read, 1, 3)
+    helper_test_combinedratelimit(burst_diff_inst.rate_limit_write, 2, 3)
+
+    class sustain_diff(RateLimitedProvider):
+        RATE_LIMITS = {"burst": 4, "sustain": {"read": 5, "write": 6}}
+
+    sustain_diff_inst = sustain_diff(xbl_client)
+
+    # Burst values are the same (second paramater)
+    helper_test_combinedratelimit(sustain_diff_inst.rate_limit_read, 4, 5)
+    helper_test_combinedratelimit(sustain_diff_inst.rate_limit_write, 4, 6)
+
+
+def test_ratelimitedprovider_rate_limits_missing_values_correct_type(xbl_client):
+    class child_class(RateLimitedProvider):
+        RATE_LIMITS = {"incorrect": "values"}
+
+    with pytest.raises(XboxException) as exception:
+        child_class(xbl_client)
+
+    ex: XboxException = exception.value
+    assert "RATE_LIMITS object missing required keys" in ex.args[0]
+
+
+def test_ratelimitedprovider_rate_limits_not_set(xbl_client):
+    class child_class(RateLimitedProvider):
+        pass
+
+    with pytest.raises(XboxException) as exception:
+        child_class(xbl_client)
+
+    ex: XboxException = exception.value
+    assert "RateLimitedProvider as parent class but RATE_LIMITS not set!" in ex.args[0]
+
+
+def test_ratelimitedprovider_rate_limits_incorrect_key_type(xbl_client):
+    class child_class(RateLimitedProvider):
+        RATE_LIMITS = {"burst": True, "sustain": False}
+
+    with pytest.raises(XboxException) as exception:
+        child_class(xbl_client)
+
+    ex: XboxException = exception.value
+    assert "RATE_LIMITS value types not recognised." in ex.args[0]
+
+
+@pytest.mark.asyncio
+async def test_ratelimits_exceeded_burst_only(respx_mock, xbl_client):
+    async def make_request():
+        route = respx_mock.get("https://social.xboxlive.com").mock(
+            return_value=Response(200, json=get_response_json("people_summary_own"))
+        )
+        ret = await xbl_client.people.get_friends_summary_own()
+
+        assert route.called
+
+    # Record the start time to ensure that the timeouts are the correct length
+    start_time = datetime.now()
+
+    # Make as many requests as possible without exceeding
+    max_request_num = xbl_client.people.RATE_LIMITS["burst"]
+    for i in range(max_request_num):
+        await make_request()
+
+    # Make another request, ensure that it raises the exception.
+    with pytest.raises(RateLimitExceededException) as exception:
+        await make_request()
+
+    # Get the error instance from pytest
+    ex: RateLimitExceededException = exception.value
+
+    # Assert that the counter matches the max request num (should not have incremented above max value)
+    assert ex.rate_limit.get_counter() == max_request_num
+
+    # Get the timeout we were issued
+    try_again_in = ex.rate_limit.get_reset_after()
+
+    # Assert that the timeout is the correct length
+    delta: timedelta = try_again_in - start_time
+    assert delta.seconds == TimePeriod.BURST.value  # 15 seconds
+
+
+async def helper_reach_and_wait_for_burst(
+    make_request, start_time, burst_limit: int, expected_counter: int
+):
+    # Make as many requests as possible without exceeding the BURST limit.
+    for i in range(burst_limit):
+        await make_request()
+
+    # Make another request, ensure that it raises the exception.
+    with pytest.raises(RateLimitExceededException) as exception:
+        await make_request()
+
+    # Get the error instance from pytest
+    ex: RateLimitExceededException = exception.value
+
+    # Assert that the counter matches the what we expect (burst, 2x burstm etc)
+    assert ex.rate_limit.get_counter() == expected_counter
+
+    # Get the reset_after value
+    # (if we call it after waiting for it to expire the function will return None)
+    burst_resets_after = ex.rate_limit.get_reset_after()
+
+    # Wait for the burst limit timeout to elapse.
+    await asyncio.sleep(TimePeriod.BURST.value)  # 15 seconds
+
+    # Assert that the reset_after value has passed.
+    assert burst_resets_after < datetime.now()
+
+
+@pytest.mark.asyncio
+async def test_ratelimits_exceeded_sustain_only(respx_mock, xbl_client):
+    async def make_request():
+        route = respx_mock.get("https://social.xboxlive.com").mock(
+            return_value=Response(200, json=get_response_json("people_summary_own"))
+        )
+        ret = await xbl_client.people.get_friends_summary_own()
+
+        assert route.called
+
+    # Record the start time to ensure that the timeouts are the correct length
+    start_time = datetime.now()
+
+    # Get the max requests for this route.
+    max_request_num = xbl_client.people.RATE_LIMITS["sustain"]  # 30
+    burst_max_request_num = xbl_client.people.RATE_LIMITS["burst"]  # 10
+
+    # In this case, the BURST limit is three times that of SUSTAIN, so we need to exceed the burst limit three times.
+
+    # Exceed the burst limit and wait for it to reset (10 requests)
+    await helper_reach_and_wait_for_burst(
+        make_request, start_time, burst_limit=burst_max_request_num, expected_counter=10
+    )
+
+    # Repeat: Exceed the burst limit and wait for it to reset (10 requests)
+    # Counter (the sustain one will be returned)
+    #         For (CombinedRateLimit).get_counter(), the highest counter is returned. (sustain in this case)
+    await helper_reach_and_wait_for_burst(
+        make_request, start_time, burst_limit=burst_max_request_num, expected_counter=20
+    )
+
+    # Now, make the rest of the requests (10 left, 20/30 done!)
+    for i in range(10):
+        await make_request()
+
+    # Wait for the burst limit to 'reset'.
+    await asyncio.sleep(TimePeriod.BURST.value)  # 15 seconds
+
+    # Now, we have made 30 requests.
+    # The counters should be as follows:
+    # - BURST: 0* (will reset on next check)
+    # - SUSTAIN: 30
+    # The next request we make should exceed the SUSTAIN rate limit.
+
+    # Make another request, ensure that it raises the exception.
+    with pytest.raises(RateLimitExceededException) as exception:
+        await make_request()
+
+    # Get the error instance from pytest
+    ex: RateLimitExceededException = exception.value
+
+    # Get the SingleRateLimit objects from the exception
+    rl: CombinedRateLimit = ex.rate_limit
+    burst = rl.get_limits_by_period(TimePeriod.BURST)[0]
+    sustain = rl.get_limits_by_period(TimePeriod.SUSTAIN)[0]
+
+    # Assert that we have only exceeded the sustain limit.
+    assert not burst.is_exceeded()
+    assert sustain.is_exceeded()
+
+    # Assert that the counter matches the max request num (should not have incremented above max value)
+    assert ex.rate_limit.get_counter() == max_request_num
+
+    # Get the timeout we were issued
+    try_again_in = ex.rate_limit.get_reset_after()
+
+    # Assert that the timeout is the correct length
+    # The SUSTAIN counter has not been reset during this test, so the try again in should be 300 seconds since we started this test.
+    delta: timedelta = try_again_in - start_time
+    assert delta.seconds == TimePeriod.SUSTAIN.value  # 300 seconds (5 minutes)

xbox/webapi/api/client.py

@@ -28,6 +28,8 @@
 from xbox.webapi.api.provider.usersearch import UserSearchProvider
 from xbox.webapi.api.provider.userstats import UserStatsProvider
 from xbox.webapi.authentication.manager import AuthenticationManager
+from xbox.webapi.common.exceptions import RateLimitExceededException
+from xbox.webapi.common.ratelimits import RateLimit
 
 log = logging.getLogger("xbox.api")
 
@@ -55,6 +57,9 @@ async def request(
         extra_params = kwargs.pop("extra_params", None)
         extra_data = kwargs.pop("extra_data", None)
 
+        # Rate limit object
+        rate_limits: RateLimit = kwargs.pop("rate_limits", None)
+
         if include_auth:
             # Ensure tokens valid
             await self._auth_mgr.refresh_tokens()
@@ -78,10 +83,20 @@ async def request(
             data = data or {}
             data.update(extra_data)
 
-        return await self._auth_mgr.session.request(
+        if rate_limits:
+            # Check if rate limits have been exceeded for this endpoint
+            if rate_limits.is_exceeded():
+                raise RateLimitExceededException("Rate limit exceeded", rate_limits)
+
+        response = await self._auth_mgr.session.request(
             method, url, **kwargs, headers=headers, params=params, data=data
         )
 
+        if rate_limits:
+            rate_limits.increment()
+
+        return response
+
     async def get(self, url: str, **kwargs: Any) -> Response:
         return await self.request("GET", url, **kwargs)
 

xbox/webapi/api/provider/achievements/__init__.py

@@ -9,14 +9,16 @@
     AchievementResponse,
     RecentProgressResponse,
 )
-from xbox.webapi.api.provider.baseprovider import BaseProvider
+from xbox.webapi.api.provider.ratelimitedprovider import RateLimitedProvider
 
 
-class AchievementsProvider(BaseProvider):
+class AchievementsProvider(RateLimitedProvider):
     ACHIEVEMENTS_URL = "https://achievements.xboxlive.com"
     HEADERS_GAME_360_PROGRESS = {"x-xbl-contract-version": "1"}
     HEADERS_GAME_PROGRESS = {"x-xbl-contract-version": "2"}
 
+    RATE_LIMITS = {"burst": 100, "sustain": 300}
+
     async def get_achievements_detail_item(
         self, xuid, service_config_id, achievement_id, **kwargs
     ) -> AchievementResponse:
@@ -33,7 +35,10 @@ async def get_achievements_detail_item(
         """
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/achievements/{service_config_id}/{achievement_id}"
         resp = await self.client.session.get(
-            url, headers=self.HEADERS_GAME_PROGRESS, **kwargs
+            url,
+            headers=self.HEADERS_GAME_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return AchievementResponse(**resp.json())
@@ -54,7 +59,11 @@ async def get_achievements_xbox360_all(
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/titleachievements?"
         params = {"titleId": title_id}
         resp = await self.client.session.get(
-            url, params=params, headers=self.HEADERS_GAME_360_PROGRESS, **kwargs
+            url,
+            params=params,
+            headers=self.HEADERS_GAME_360_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return Achievement360Response(**resp.json())
@@ -75,7 +84,11 @@ async def get_achievements_xbox360_earned(
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/achievements?"
         params = {"titleId": title_id}
         resp = await self.client.session.get(
-            url, params=params, headers=self.HEADERS_GAME_360_PROGRESS, **kwargs
+            url,
+            params=params,
+            headers=self.HEADERS_GAME_360_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return Achievement360Response(**resp.json())
@@ -94,7 +107,10 @@ async def get_achievements_xbox360_recent_progress_and_info(
         """
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/history/titles"
         resp = await self.client.session.get(
-            url, headers=self.HEADERS_GAME_360_PROGRESS, **kwargs
+            url,
+            headers=self.HEADERS_GAME_360_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return Achievement360ProgressResponse(**resp.json())
@@ -115,7 +131,11 @@ async def get_achievements_xboxone_gameprogress(
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/achievements?"
         params = {"titleId": title_id}
         resp = await self.client.session.get(
-            url, params=params, headers=self.HEADERS_GAME_PROGRESS, **kwargs
+            url,
+            params=params,
+            headers=self.HEADERS_GAME_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return AchievementResponse(**resp.json())
@@ -134,7 +154,10 @@ async def get_achievements_xboxone_recent_progress_and_info(
         """
         url = f"{self.ACHIEVEMENTS_URL}/users/xuid({xuid})/history/titles"
         resp = await self.client.session.get(
-            url, headers=self.HEADERS_GAME_PROGRESS, **kwargs
+            url,
+            headers=self.HEADERS_GAME_PROGRESS,
+            rate_limits=self.rate_limit_read,
+            **kwargs,
         )
         resp.raise_for_status()
         return RecentProgressResponse(**resp.json())