From 2a3cdbffa6ace8c0966452a4185623bc39fa32df Mon Sep 17 00:00:00 2001
From: "Diego M. Rodriguez" <diego.plan9@gmail.com>
Date: Mon, 4 Oct 2021 09:38:29 +0200
Subject: [PATCH] Support pyjwt > 2 in tests

Allow the tests to be aware of the `pyjwt` version installed, converting
the result of `jwt.encode()` from bytes to string and vice-versa
accordingly as the return type changes between pyjwt versions.
---
 oauthenticator/tests/test_azuread.py   | 5 ++++-
 oauthenticator/tests/test_mediawiki.py | 7 ++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/oauthenticator/tests/test_azuread.py b/oauthenticator/tests/test_azuread.py
index 35c8a2bf..90ce1c1b 100644
--- a/oauthenticator/tests/test_azuread.py
+++ b/oauthenticator/tests/test_azuread.py
@@ -9,6 +9,7 @@
 import pytest
 
 from ..azuread import AzureAdOAuthenticator
+from ..azuread import PYJWT_2
 from .mocks import setup_oauth_mock
 
 
@@ -40,7 +41,9 @@ def user_model(tenant_id, client_id, name):
             "aio": "Df2UVXL1ix!lMCWMSOJBcFatzcGfvFGhjKv8q5g0x732dR5MB5BisvGQO7YWByjd8iQDLq!eGbIDakyp5mnOrcdqHeYSnltepQmRp6AIZ8jY",
         },
         os.urandom(5),
-    ).decode("ascii")
+    )
+    if not PYJWT_2:
+        id_token = id_token.decode("ascii")
 
     return {
         "access_token": "abc123",
diff --git a/oauthenticator/tests/test_mediawiki.py b/oauthenticator/tests/test_mediawiki.py
index 8a813d06..6b9cfa7c 100644
--- a/oauthenticator/tests/test_mediawiki.py
+++ b/oauthenticator/tests/test_mediawiki.py
@@ -8,6 +8,7 @@
 from pytest import fixture
 from tornado import web
 
+from ..azuread import PYJWT_2
 from ..mediawiki import AUTH_REQUEST_COOKIE_NAME
 from ..mediawiki import MWOAuthenticator
 from .mocks import mock_handler
@@ -20,7 +21,7 @@ def mediawiki():
     def post_token(request, context):
         authorization_header = request.headers['Authorization'].decode('utf8')
         request_nonce = re.search(r'oauth_nonce="(.*?)"', authorization_header).group(1)
-        return jwt.encode(
+        content = jwt.encode(
             {
                 'username': 'wash',
                 'aud': 'client_id',
@@ -30,6 +31,10 @@ def post_token(request, context):
             },
             'client_secret',
         )
+        if PYJWT_2:
+            content = content.encode()
+
+        return content
 
     with requests_mock.Mocker() as mock:
         mock.post(