Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/client/auth_token_loader.py: 100.00%

1# Copyright (C) 2023 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15from cachetools import TTLCache 

16 

17from buildgrid.server.exceptions import InvalidArgumentError 

18 

19DEFAULT_TOKEN_REFRESH_SECONDS = 60 * 60 

20_CACHE_KEY = "token" 

21 

22 

23class AuthTokenLoader: 

24 

25 def __init__(self, token_path: str, refresh_in_seconds: int | None = None) -> None: 

26 self._token_path = token_path 

27 self._cache: TTLCache[str, str] = TTLCache( 

28 maxsize=1, ttl=refresh_in_seconds if refresh_in_seconds else DEFAULT_TOKEN_REFRESH_SECONDS 

29 ) 

30 

31 def _load_token_from_file(self) -> None: 

32 try: 

33 with open(self._token_path, mode="r") as token_file: 

34 token = token_file.read().strip() 

35 self._cache[_CACHE_KEY] = token 

36 except Exception as e: 

37 raise InvalidArgumentError(f"Cannot read token from filepath: {self._token_path}") from e 

38 

39 def get_token(self) -> str: 

40 if _CACHE_KEY not in self._cache: 

41 self._load_token_from_file() 

42 

43 assert _CACHE_KEY in self._cache 

44 return self._cache[_CACHE_KEY]