diff --git a/src/kili/core/graphql/graphql_client.py b/src/kili/core/graphql/graphql_client.py index b4947acc1..b63635172 100644 --- a/src/kili/core/graphql/graphql_client.py +++ b/src/kili/core/graphql/graphql_client.py @@ -3,6 +3,7 @@ import logging import os import threading +import time from pathlib import Path from typing import Any, Dict, Optional, Union from urllib.parse import urlparse @@ -80,6 +81,8 @@ def __init__( self.http_client = http_client self.verify = verify self.enable_schema_caching = enable_schema_caching + self.created_at = time.time() + self.complexity_consumed = 0 self.graphql_schema_cache_dir = ( Path(graphql_schema_cache_dir) if graphql_schema_cache_dir else None ) @@ -309,7 +312,7 @@ def _raw_execute( ) -> Dict[str, Any]: _limiter.try_acquire("GraphQLClient.execute") with _execute_lock: - return self._gql_client.execute( + res = self._gql_client.execute( document=document, variable_values=variables, extra_args={ @@ -320,3 +323,9 @@ def _raw_execute( }, **kwargs, ) + transport = self._gql_client.transport + if transport: + headers = transport.response_headers # pyright: ignore[reportGeneralTypeIssues] + returned_complexity = int(headers.get("x-complexity", 0)) if headers else 0 + self.complexity_consumed += returned_complexity + return res diff --git a/tests/integration/entrypoints/client/test_client.py b/tests/integration/entrypoints/client/test_client.py index 08f52b3f4..5c1658ba1 100644 --- a/tests/integration/entrypoints/client/test_client.py +++ b/tests/integration/entrypoints/client/test_client.py @@ -140,3 +140,28 @@ def test_given_an_api_key_away_to_expiration_when_I_check_expiry_of_key_is_not_c warnings.simplefilter("error") # checks that no warning is raised # When _ = Kili() + + +@patch.dict(os.environ, {"KILI_API_KEY": "fake_key"}) +def test_complexity_increases_with_calls( + mocker: pytest_mock.MockerFixture, +): + graphql_mock = mocker.MagicMock() + graphql_mock.execute.return_value = {"data": 1} + graphql_mock.transport.response_headers = {"x-complexity": "125"} + + # Given + mocker.patch("kili.client.is_api_key_valid", return_value=True) + mocker.patch.object(ApiKeyUseCases, "check_expiry_of_key_is_close") + mocker.patch( + "kili.core.graphql.graphql_client.GraphQLClient._initizalize_graphql_client", + return_value=graphql_mock, + ) + kili = Kili() + + assert kili.graphql_client.complexity_consumed == 0 + initialized_date = kili.graphql_client.created_at + + kili.count_projects() + assert kili.graphql_client.complexity_consumed == 125 + assert kili.graphql_client.created_at == initialized_date