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 typing import Optional 

16 

17from cachetools import TTLCache 

18 

19from buildgrid.server.exceptions import InvalidArgumentError 

20 

21DEFAULT_TOKEN_REFRESH_SECONDS = 60 * 60 

22_CACHE_KEY = "token" 

23 

24 

25class AuthTokenLoader: 

26 

27 def __init__(self, token_path: str, refresh_in_seconds: Optional[int] = None) -> None: 

28 self._token_path = token_path 

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

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

31 ) 

32 

33 def _load_token_from_file(self) -> None: 

34 try: 

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

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

37 self._cache[_CACHE_KEY] = token 

38 except Exception as e: 

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

40 

41 def get_token(self) -> str: 

42 if _CACHE_KEY not in self._cache: 

43 self._load_token_from_file() 

44 

45 assert _CACHE_KEY in self._cache 

46 return self._cache[_CACHE_KEY]