Repository: jpadilla/pyjwt
Pull request #781 (2022-07-11 06:27:23)

Add cacheing functionality for JWK set
## Problem
The current PyJWT library only support caching of single signing key. Ref to https://github.com/jpadilla/pyjwt/pull/611

Even though it avoids the network call when the services periodically calling `get_signing_key` with same `kid`, the current implementation doesn't have any TTL, so the cache is stored forever. Also, it still makes network call when different `kid` is provided.

## Solution & Implantation
- Add `JWKSetCache` object that stores jwk set with current timestamp
    - When user try to get signing key from jwk set, the service will try to check if the cache exist and not expired. If so, it will avoid the network call and return the cached jwk set.
    - If the network call throws any error, clear the cache.
    - If the service can't find matching `kid` from the jwk set, it will make a new network call to get a new jwk set and check again in new set. Raise error if still cannot find the signing key with matching kid. 
    - Add option to enable/disable jwk set cache and lifespan. Default lifespan is 5 minutes. 
- Add unit tests to cover the added code
- Disable LRU cache for `get_signing_key` by default as suggested from this comment https://github.com/jpadilla/pyjwt/pull/611#issuecomment-815807580


Patch:
diff --git a/docs/api.rst b/docs/api.rst
index 2f81b1f..919b6af 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -62,7 +62,7 @@ API Reference
             if ``verify_exp``, ``verify_iat``, and ``verify_nbf`` respectively
             is set to ``True``).
 
-    :param Iterable audience: optional, the value for ``verify_aud`` check
+    :param Union[str, Iterable] audience: optional, the value for ``verify_aud`` check
     :param str issuer: optional, the value for ``verify_iss`` check
     :param float leeway: a time margin in seconds for the expiration check
     :rtype: dict
diff --git a/jwt/api_jwk.py b/jwt/api_jwk.py
index aae2cf1..aa3dd32 100644
--- a/jwt/api_jwk.py
+++ b/jwt/api_jwk.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import json
+import time
 
 from .algorithms import get_default_algorithms
 from .exceptions import InvalidKeyError, PyJWKError, PyJWKSetError
@@ -74,15 +77,15 @@ class PyJWK:
 
 
 class PyJWKSet:
-    def __init__(self, keys):
+    def __init__(self, keys: list[dict]) -> None:
         self.keys = []
 
-        if not keys or not isinstance(keys, list):
-            raise PyJWKSetError("Invalid JWK Set value")
-
-        if len(keys) == 0:
+        if not keys:
             raise PyJWKSetError("The JWK Set did not contain any keys")
 
+        if not isinstance(keys, list):
+            raise PyJWKSetError("Invalid JWK Set value")
+
         for key in keys:
             try:
                 self.keys.append(PyJWK(key))
@@ -108,3 +111,15 @@ class PyJWKSet:
             if key.key_id == kid:
                 return key
         raise KeyError(f"keyset has no key for kid: {kid}")
+
+
+class PyJWTSetWithTimestamp:
+    def __init__(self, jwk_set: PyJWKSet):
+        self.jwk_set = jwk_set
+        self.timestamp = time.monotonic()
+
+    def get_jwk_set(self):
+        return self.jwk_set
+
+    def get_timestamp(self):
+        return self.timestamp
diff --git a/jwt/api_jwt.py b/jwt/api_jwt.py
index b08e950..91a6d2e 100644
--- a/jwt/api_jwt.py
+++ b/jwt/api_jwt.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import warnings
 from calendar import timegm
@@ -76,7 +78,7 @@ class PyJWT:
         detached_payload: Optional[bytes] = None,
         # passthrough arguments to _validate_claims
         # consider putting in options
-        audience: Optional[str] = None,
+        audience: Optional[Union[str, Iterable[str]]] = None,
         issuer: Optional[str] = None,
         leeway: Union[int, float, timedelta] = 0,
         # kwargs
@@ -150,7 +152,7 @@ class PyJWT:
         detached_payload: Optional[bytes] = None,
         # passthrough arguments to _validate_claims
         # consider putting in options
-        audience: Optional[str] = None,
+        audience: Optional[Union[str, Iterable[str]]] = None,
         issuer: Optional[str] = None,
         leeway: Union[int, float, timedelta] = 0,
         # kwargs
@@ -180,8 +182,8 @@ class PyJWT:
         if isinstance(leeway, timedelta):
             leeway = leeway.total_seconds()
 
-        if not isinstance(audience, (bytes, str, type(None), Iterable)):
-            raise TypeError("audience must be a string, iterable, or None")
+        if audience is not None and not isinstance(audience, (str, Iterable)):
+            raise TypeError("audience must be a string, iterable or None")
 
         self._validate_required_claims(payload, options)
 
diff --git a/jwt/help.py b/jwt/help.py
index d5c3ebb..767c323 100644
--- a/jwt/help.py
+++ b/jwt/help.py
@@ -1,6 +1,7 @@
 import json
 import platform
 import sys
+from typing import Dict
 
 from . import __version__ as pyjwt_version
 
@@ -10,7 +11,7 @@ except ModuleNotFoundError:
     cryptography = None  # type: ignore
 
 
-def info():
+def info() -> Dict[str, Dict[str, str]]:
     """
     Generate information for a bug report.
     Based on the requests package help utility module.
@@ -29,13 +30,13 @@ def info():
         implementation_version = platform.python_version()
     elif implementation == "PyPy":
         implementation_version = (
-            f"{sys.pypy_version_info.major}."
+            f"{sys.pypy_version_info.major}."  # type: ignore[attr-defined]
             f"{sys.pypy_version_info.minor}."
             f"{sys.pypy_version_info.micro}"
         )
-        if sys.pypy_version_info.releaselevel != "final":
+        if sys.pypy_version_info.releaselevel != "final":  # type: ignore[attr-defined]
             implementation_version = "".join(
-                [implementation_version, sys.pypy_version_info.releaselevel]
+                [implementation_version, sys.pypy_version_info.releaselevel]  # type: ignore[attr-defined]
             )
     else:
         implementation_version = "Unknown"
@@ -51,7 +52,7 @@ def info():
     }
 
 
-def main():
+def main() -> None:
     """Pretty-print the bug information as JSON."""
     print(json.dumps(info(), sort_keys=True, indent=2))
 
diff --git a/jwt/jwk_set_cache.py b/jwt/jwk_set_cache.py
new file mode 100644
index 0000000..e8c2a7e
--- /dev/null
+++ b/jwt/jwk_set_cache.py
@@ -0,0 +1,32 @@
+import time
+from typing import Optional
+
+from .api_jwk import PyJWKSet, PyJWTSetWithTimestamp
+
+
+class JWKSetCache:
+    def __init__(self, lifespan: int):
+        self.jwk_set_with_timestamp: Optional[PyJWTSetWithTimestamp] = None
+        self.lifespan = lifespan
+
+    def put(self, jwk_set: PyJWKSet):
+        if jwk_set is not None:
+            self.jwk_set_with_timestamp = PyJWTSetWithTimestamp(jwk_set)
+        else:
+            # clear cache
+            self.jwk_set_with_timestamp = None
+
+    def get(self) -> Optional[PyJWKSet]:
+        if self.jwk_set_with_timestamp is None or self.is_expired():
+            return None
+
+        return self.jwk_set_with_timestamp.get_jwk_set()
+
+    def is_expired(self) -> bool:
+
+        return (
+            self.jwk_set_with_timestamp is not None
+            and self.lifespan > -1
+            and time.monotonic()
+            > self.jwk_set_with_timestamp.get_timestamp() + self.lifespan
+        )
diff --git a/jwt/jwks_client.py b/jwt/jwks_client.py
index 767b717..b4e9800 100644
--- a/jwt/jwks_client.py
+++ b/jwt/jwks_client.py
@@ -1,31 +1,68 @@
 import json
 import urllib.request
 from functools import lru_cache
-from typing import Any, List
+from typing import Any, List, Optional
+from urllib.error import URLError
 
 from .api_jwk import PyJWK, PyJWKSet
 from .api_jwt import decode_complete as decode_token
 from .exceptions import PyJWKClientError
+from .jwk_set_cache import JWKSetCache
 
 
 class PyJWKClient:
-    def __init__(self, uri: str, cache_keys: bool = True, max_cached_keys: int = 16):
+    def __init__(
+        self,
+        uri: str,
+        cache_keys: bool = False,
+        max_cached_keys: int = 16,
+        cache_jwk_set: bool = True,
+        lifespan: int = 300,
+    ):
         self.uri = uri
+        self.jwk_set_cache: Optional[JWKSetCache] = None
+
+        if cache_jwk_set:
+            # Init jwt set cache with default or given lifespan.
+            # Default lifespan is 300 seconds (5 minutes).
+            if lifespan <= 0:
+                raise PyJWKClientError(
+                    f'Lifespan must be greater than 0, the input is "{lifespan}"'
+                )
+            self.jwk_set_cache = JWKSetCache(lifespan)
+        else:
+            self.jwk_set_cache = None
+
         if cache_keys:
             # Cache signing keys
             # Ignore mypy (https://github.com/python/mypy/issues/2427)
             self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key)  # type: ignore
 
     def fetch_data(self) -> Any:
-        with urllib.request.urlopen(self.uri) as response:
-            return json.load(response)
+        jwk_set: Any = None
+        try:
+            with urllib.request.urlopen(self.uri) as response:
+                jwk_set = json.load(response)
+        except URLError as e:
+            raise PyJWKClientError(f'Fail to fetch data from the url, err: "{e}"')
+        else:
+            return jwk_set
+        finally:
+            if self.jwk_set_cache is not None:
+                self.jwk_set_cache.put(jwk_set)
+
+    def get_jwk_set(self, refresh: bool = False) -> PyJWKSet:
+        data = None
+        if self.jwk_set_cache is not None and not refresh:
+            data = self.jwk_set_cache.get()
+
+        if data is None:
+            data = self.fetch_data()
 
-    def get_jwk_set(self) -> PyJWKSet:
-        data = self.fetch_data()
         return PyJWKSet.from_dict(data)
 
-    def get_signing_keys(self) -> List[PyJWK]:
-        jwk_set = self.get_jwk_set()
+    def get_signing_keys(self, refresh: bool = False) -> List[PyJWK]:
+        jwk_set = self.get_jwk_set(refresh)
         signing_keys = [
             jwk_set_key
             for jwk_set_key in jwk_set.keys
@@ -39,17 +76,17 @@ class PyJWKClient:
 
     def get_signing_key(self, kid: str) -> PyJWK:
         signing_keys = self.get_signing_keys()
-        signing_key = None
-
-        for key in signing_keys:
-            if key.key_id == kid:
-                signing_key = key
-                break
+        signing_key = self.match_kid(signing_keys, kid)
 
         if not signing_key:
-            raise PyJWKClientError(
-                f'Unable to find a signing key that matches: "{kid}"'
-            )
+            # If no matching signing key from the jwk set, refresh the jwk set and try again.
+            signing_keys = self.get_signing_keys(refresh=True)
+            signing_key = self.match_kid(signing_keys, kid)
+
+            if not signing_key:
+                raise PyJWKClientError(
+                    f'Unable to find a signing key that matches: "{kid}"'
+                )
 
         return signing_key
 
@@ -57,3 +94,14 @@ class PyJWKClient:
         unverified = decode_token(token, options={"verify_signature": False})
         header = unverified["header"]
         return self.get_signing_key(header.get("kid"))
+
+    @staticmethod
+    def match_kid(signing_keys: List[PyJWK], kid: str) -> Optional[PyJWK]:
+        signing_key = None
+
+        for key in signing_keys:
+            if key.key_id == kid:
+                signing_key = key
+                break
+
+        return signing_key
diff --git a/setup.cfg b/setup.cfg
index 67d30a7..9fc9169 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -45,6 +45,7 @@ docs =
     zope.interface
 crypto =
     cryptography>=3.3.1
+    types-cryptography>=3.3.21
 tests =
     pytest>=6.0.0,<7.0.0
     coverage[toml]==5.0.4
@@ -53,6 +54,7 @@ dev =
     sphinx-rtd-theme
     zope.interface
     cryptography>=3.3.1
+    types-cryptography>=3.3.21
     pytest>=6.0.0,<7.0.0
     coverage[toml]==5.0.4
     mypy


Test patch:
diff --git a/tests/test_api_jwk.py b/tests/test_api_jwk.py
index 94eb498..040e81b 100644
--- a/tests/test_api_jwk.py
+++ b/tests/test_api_jwk.py
@@ -208,8 +208,8 @@ class TestPyJWK:
             PyJWK.from_dict(v)
 
 
+@crypto_required
 class TestPyJWKSet:
-    @crypto_required
     def test_should_load_keys_from_jwk_data_dict(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
 
@@ -231,7 +231,6 @@ class TestPyJWKSet:
         assert jwk.key_id == "keyid-abc123"
         assert jwk.public_key_use == "sig"
 
-    @crypto_required
     def test_should_load_keys_from_jwk_data_json_string(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
 
@@ -253,7 +252,6 @@ class TestPyJWKSet:
         assert jwk.key_id == "keyid-abc123"
         assert jwk.public_key_use == "sig"
 
-    @crypto_required
     def test_keyset_should_index_by_kid(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
 
@@ -276,7 +274,6 @@ class TestPyJWKSet:
         with pytest.raises(KeyError):
             _ = jwk_set["this-kid-does-not-exist"]
 
-    @crypto_required
     def test_keyset_with_unknown_alg(self):
         # first keyset with unusable key and usable key
         with open(key_path("jwk_keyset_with_unknown_alg.json")) as keyfile:
@@ -293,3 +290,13 @@ class TestPyJWKSet:
             assert len(jwks.get("keys")) == 1
             with pytest.raises(PyJWKSetError):
                 _ = PyJWKSet.from_json(jwks_text)
+
+    def test_invalid_keys_list(self):
+        with pytest.raises(PyJWKSetError) as err:
+            PyJWKSet(keys="string")
+        assert str(err.value) == "Invalid JWK Set value"
+
+    def test_empty_keys_list(self):
+        with pytest.raises(PyJWKSetError) as err:
+            PyJWKSet(keys=[])
+        assert str(err.value) == "The JWK Set did not contain any keys"
diff --git a/tests/test_api_jwt.py b/tests/test_api_jwt.py
index d0443e8..bebe7d2 100644
--- a/tests/test_api_jwt.py
+++ b/tests/test_api_jwt.py
@@ -119,7 +119,7 @@ class TestJWT:
             jwt.decode(example_jwt, secret, audience=1, algorithms=["HS256"])
 
         exception = context.value
-        assert str(exception) == "audience must be a string, iterable, or None"
+        assert str(exception) == "audience must be a string, iterable or None"
 
     def test_decode_with_nonlist_aud_claim_throws_exception(self, jwt):
         secret = "secret"
@@ -419,6 +419,14 @@ class TestJWT:
         with pytest.raises(InvalidAudienceError):
             jwt.decode(token, "secret", audience="urn-me", algorithms=["HS256"])
 
+    def test_raise_exception_audience_as_bytes(self, jwt):
+        payload = {"some": "payload", "aud": ["urn:me", "urn:someone-else"]}
+        token = jwt.encode(payload, "secret")
+        with pytest.raises(InvalidAudienceError):
+            jwt.decode(
+                token, "secret", audience="urn:me".encode(), algorithms=["HS256"]
+            )
+
     def test_raise_exception_invalid_audience_in_array(self, jwt):
         payload = {
             "some": "payload",
diff --git a/tests/test_jwks_client.py b/tests/test_jwks_client.py
index 3e42da1..c95dfcc 100644
--- a/tests/test_jwks_client.py
+++ b/tests/test_jwks_client.py
@@ -1,6 +1,8 @@
 import contextlib
 import json
+import time
 from unittest import mock
+from urllib.error import URLError
 
 import pytest
 
@@ -11,7 +13,7 @@ from jwt.exceptions import PyJWKClientError
 
 from .utils import crypto_required
 
-RESPONSE_DATA = {
+RESPONSE_DATA_WITH_MATCHING_KID = {
     "keys": [
         {
             "alg": "RS256",
@@ -28,9 +30,22 @@ RESPONSE_DATA = {
     ]
 }
 
+RESPONSE_DATA_NO_MATCHING_KID = {
+    "keys": [
+        {
+            "alg": "RS256",
+            "kty": "RSA",
+            "use": "sig",
+            "n": "39SJ39VgrQ0qMNK74CaueUBlyYsUyuA7yWlHYZ-jAj6tlFKugEVUTBUVbhGF44uOr99iL_cwmr-srqQDEi-jFHdkS6WFkYyZ03oyyx5dtBMtzrXPieFipSGfQ5EGUGloaKDjL-Ry9tiLnysH2VVWZ5WDDN-DGHxuCOWWjiBNcTmGfnj5_NvRHNUh2iTLuiJpHbGcPzWc5-lc4r-_ehw9EFfp2XsxE9xvtbMZ4SouJCiv9xnrnhe2bdpWuu34hXZCrQwE8DjRY3UR8LjyMxHHPLzX2LWNMHjfN3nAZMteS-Ok11VYDFI-4qCCVGo_WesBCAeqCjPLRyZoV27x1YGsUQ",
+            "e": "AQAB",
+            "kid": "MLYHNMMhwCNXw9roHIILFsK4nLs=",
+        }
+    ]
+}
+
 
 @contextlib.contextmanager
-def mocked_response(data):
+def mocked_success_response(data):
     with mock.patch("urllib.request.urlopen") as urlopen_mock:
         response = mock.Mock()
         response.__enter__ = mock.Mock(return_value=response)
@@ -40,12 +55,35 @@ def mocked_response(data):
         yield urlopen_mock
 
 
+@contextlib.contextmanager
+def mocked_failed_response():
+    with mock.patch("urllib.request.urlopen") as urlopen_mock:
+        urlopen_mock.side_effect = URLError("Fail to process the request.")
+        yield urlopen_mock
+
+
+@contextlib.contextmanager
+def mocked_first_call_wrong_kid_second_call_correct_kid(
+    response_data_one, response_data_two
+):
+    with mock.patch("urllib.request.urlopen") as urlopen_mock:
+        response = mock.Mock()
+        response.__enter__ = mock.Mock(return_value=response)
+        response.__exit__ = mock.Mock()
+        response.read.side_effect = [
+            json.dumps(response_data_one),
+            json.dumps(response_data_two),
+        ]
+        urlopen_mock.return_value = response
+        yield urlopen_mock
+
+
 @crypto_required
 class TestPyJWKClient:
     def test_get_jwk_set(self):
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client = PyJWKClient(url)
             jwk_set = jwks_client.get_jwk_set()
 
@@ -54,7 +92,7 @@ class TestPyJWKClient:
     def test_get_signing_keys(self):
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client = PyJWKClient(url)
             signing_keys = jwks_client.get_signing_keys()
 
@@ -64,11 +102,11 @@ class TestPyJWKClient:
     def test_get_signing_keys_if_no_use_provided(self):
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
 
-        mocked_key = RESPONSE_DATA["keys"][0].copy()
+        mocked_key = RESPONSE_DATA_WITH_MATCHING_KID["keys"][0].copy()
         del mocked_key["use"]
         response = {"keys": [mocked_key]}
 
-        with mocked_response(response):
+        with mocked_success_response(response):
             jwks_client = PyJWKClient(url)
             signing_keys = jwks_client.get_signing_keys()
 
@@ -78,10 +116,10 @@ class TestPyJWKClient:
     def test_get_signing_keys_raises_if_none_found(self):
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
 
-        mocked_key = RESPONSE_DATA["keys"][0].copy()
+        mocked_key = RESPONSE_DATA_WITH_MATCHING_KID["keys"][0].copy()
         mocked_key["use"] = "enc"
         response = {"keys": [mocked_key]}
-        with mocked_response(response):
+        with mocked_success_response(response):
             jwks_client = PyJWKClient(url)
 
             with pytest.raises(PyJWKClientError) as exc:
@@ -93,7 +131,7 @@ class TestPyJWKClient:
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
         kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client = PyJWKClient(url)
             signing_key = jwks_client.get_signing_key(kid)
 
@@ -106,14 +144,14 @@ class TestPyJWKClient:
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
         kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
 
-        jwks_client = PyJWKClient(url)
+        jwks_client = PyJWKClient(url, cache_keys=True)
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client.get_signing_key(kid)
 
         # mocked_response does not allow urllib.request.urlopen to be called twice
         # so a second mock is needed
-        with mocked_response(RESPONSE_DATA) as repeated_call:
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID) as repeated_call:
             jwks_client.get_signing_key(kid)
 
         assert repeated_call.call_count == 0
@@ -122,14 +160,14 @@ class TestPyJWKClient:
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
         kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
 
-        jwks_client = PyJWKClient(url, cache_keys=False)
+        jwks_client = PyJWKClient(url, cache_jwk_set=False)
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client.get_signing_key(kid)
 
         # mocked_response does not allow urllib.request.urlopen to be called twice
         # so a second mock is needed
-        with mocked_response(RESPONSE_DATA) as repeated_call:
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID) as repeated_call:
             jwks_client.get_signing_key(kid)
 
         assert repeated_call.call_count == 1
@@ -138,7 +176,7 @@ class TestPyJWKClient:
         token = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Ik5FRTFRVVJCT1RNNE16STVSa0ZETlRZeE9UVTFNRGcyT0Rnd1EwVXpNVGsxUWpZeVJrUkZRdyJ9.eyJpc3MiOiJodHRwczovL2Rldi04N2V2eDlydS5hdXRoMC5jb20vIiwic3ViIjoiYVc0Q2NhNzl4UmVMV1V6MGFFMkg2a0QwTzNjWEJWdENAY2xpZW50cyIsImF1ZCI6Imh0dHBzOi8vZXhwZW5zZXMtYXBpIiwiaWF0IjoxNTcyMDA2OTU0LCJleHAiOjE1NzIwMDY5NjQsImF6cCI6ImFXNENjYTc5eFJlTFdVejBhRTJINmtEME8zY1hCVnRDIiwiZ3R5IjoiY2xpZW50LWNyZWRlbnRpYWxzIn0.PUxE7xn52aTCohGiWoSdMBZGiYAHwE5FYie0Y1qUT68IHSTXwXVd6hn02HTah6epvHHVKA2FqcFZ4GGv5VTHEvYpeggiiZMgbxFrmTEY0csL6VNkX1eaJGcuehwQCRBKRLL3zKmA5IKGy5GeUnIbpPHLHDxr-GXvgFzsdsyWlVQvPX2xjeaQ217r2PtxDeqjlf66UYl6oY6AqNS8DH3iryCvIfCcybRZkc_hdy-6ZMoKT6Piijvk_aXdm7-QQqKJFHLuEqrVSOuBqqiNfVrG27QzAPuPOxvfXTVLXL2jek5meH6n-VWgrBdoMFH93QEszEDowDAEhQPHVs0xj7SIzA"
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
 
-        with mocked_response(RESPONSE_DATA):
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
             jwks_client = PyJWKClient(url)
             signing_key = jwks_client.get_signing_key_from_jwt(token)
 
@@ -159,3 +197,102 @@ class TestPyJWKClient:
             "azp": "aW4Cca79xReLWUz0aE2H6kD0O3cXBVtC",
             "gty": "client-credentials",
         }
+
+    def test_get_jwk_set_caches_result(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        jwks_client = PyJWKClient(url)
+        assert jwks_client.jwk_set_cache is not None
+
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
+            jwks_client.get_jwk_set()
+
+        # mocked_response does not allow urllib.request.urlopen to be called twice
+        # so a second mock is needed
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID) as repeated_call:
+            jwks_client.get_jwk_set()
+
+        assert repeated_call.call_count == 0
+
+    def test_get_jwt_set_cache_expired_result(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        jwks_client = PyJWKClient(url, lifespan=1)
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
+            jwks_client.get_jwk_set()
+
+        time.sleep(2)
+
+        # mocked_response does not allow urllib.request.urlopen to be called twice
+        # so a second mock is needed
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID) as repeated_call:
+            jwks_client.get_jwk_set()
+
+        assert repeated_call.call_count == 1
+
+    def test_get_jwt_set_cache_disabled(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        jwks_client = PyJWKClient(url, cache_jwk_set=False)
+        assert jwks_client.jwk_set_cache is None
+
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
+            jwks_client.get_jwk_set()
+
+        assert jwks_client.jwk_set_cache is None
+
+        time.sleep(2)
+
+        # mocked_response does not allow urllib.request.urlopen to be called twice
+        # so a second mock is needed
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID) as repeated_call:
+            jwks_client.get_jwk_set()
+
+        assert repeated_call.call_count == 1
+
+    def test_get_jwt_set_failed_request_should_clear_cache(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        jwks_client = PyJWKClient(url)
+        with mocked_success_response(RESPONSE_DATA_WITH_MATCHING_KID):
+            jwks_client.get_jwk_set()
+
+        with pytest.raises(PyJWKClientError):
+            with mocked_failed_response():
+                jwks_client.get_jwk_set(refresh=True)
+
+            assert jwks_client.jwk_set_cache is None
+
+    def test_get_jwt_set_refresh_cache(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        jwks_client = PyJWKClient(url)
+
+        kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
+
+        # The first call will return response with no matching kid,
+        # the function should make another call to try to refresh the cache.
+        with mocked_first_call_wrong_kid_second_call_correct_kid(
+            RESPONSE_DATA_NO_MATCHING_KID, RESPONSE_DATA_WITH_MATCHING_KID
+        ) as call_data:
+            jwks_client.get_signing_key(kid)
+
+        assert call_data.call_count == 2
+
+    def test_get_jwt_set_no_matching_kid_after_second_attempt(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        jwks_client = PyJWKClient(url)
+
+        kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
+
+        with pytest.raises(PyJWKClientError):
+            with mocked_first_call_wrong_kid_second_call_correct_kid(
+                RESPONSE_DATA_NO_MATCHING_KID, RESPONSE_DATA_NO_MATCHING_KID
+            ):
+                jwks_client.get_signing_key(kid)
+
+    def test_get_jwt_set_invalid_lifespan(self):
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        with pytest.raises(PyJWKClientError):
+            jwks_client = PyJWKClient(url, lifespan=-1)
+            assert jwks_client is None