use traitlet.default for username_claim

oauthenticator/azuread.py

@@ -6,23 +6,16 @@
 import json
 import jwt
 import os
-import re
-import string
 import urllib
-import sys
 
 from tornado.auth import OAuth2Mixin
 from tornado.log import app_log
-from tornado import web
-
-from tornado.httputil import url_concat
 from tornado.httpclient import HTTPRequest, AsyncHTTPClient
 
 from jupyterhub.auth import LocalAuthenticator
 
-from traitlets import List, Set, Unicode
+from traitlets import Unicode, default
 
-from .common import next_page_from_links
 from .oauth2 import OAuthLoginHandler, OAuthenticator
 
 
@@ -65,14 +58,8 @@ def get_tenant(self):
             app_log.info('ID4: {0}'.format(tenant_id))
             return tenant_id
 
-    def get_username_claim(self):
-        """
-        The claim to map to the jupyter username, such as `upn` or `unique_name`
-        See https://docs.microsoft.com/en-gb/azure/active-directory/develop/id-tokens
-        """
-        if hasattr(self, 'username_claim') and self.username_claim:
-            app_log.info('ID5: {0}'.format(self.username_claim))
-            return self.username_claim
+    @default('username_claim')
+    def _username_claim_default(self):
         return 'name'
 
     async def authenticate(self, handler, data=None):
@@ -112,7 +99,7 @@ async def authenticate(self, handler, data=None):
         id_token = resp_json['id_token']
         decoded = jwt.decode(id_token, verify=False)
 
-        userdict = {"name": decoded[self.get_username_claim()]}
+        userdict = {"name": decoded[self.username_claim]}
         userdict["auth_state"] = auth_state = {}
         auth_state['access_token'] = access_token
         # results in a decoded JWT for the user data

oauthenticator/tests/test_azuread.py

@@ -2,12 +2,10 @@
 from ..azuread import AzureAdOAuthenticator
 
 _t_id = 'XXX-XXX-XXXX'
-_t_username_claim = 'upn'
 
 
 class Config(object):
     tenant_id = _t_id
-    username_claim = _t_username_claim
 
 
 def test_gettenant_with_tenant_id():
@@ -21,18 +19,4 @@ def test_gettenant_with_tenant_id():
 
 def test_gettenant_from_env():
     t_id = AzureAdOAuthenticator.get_tenant(object)
-    assert t_id.default_value == "some_random_id"
-
-
-def test_username_claim_config():
-    t_username_claim = AzureAdOAuthenticator.get_username_claim(Config())
-    assert t_username_claim == _t_username_claim
-
-
-def test_username_claim_default():
-
-    class Config(object):
-        tenant_id = _t_id
-
-    t_username_claim = AzureAdOAuthenticator.get_username_claim(Config())
-    assert t_username_claim == 'name'
+    assert t_id.default_value == "some_random_id"