From 43f59a21fc2f43fa74fc6c3b503515e1771198c8 Mon Sep 17 00:00:00 2001
From: Sam Bull <git@sambull.org>
Date: Sat, 19 Dec 2020 16:44:41 +0000
Subject: [PATCH] Move indexing back into .identify() method.

---
 aiohttp_security/jwt_identity.py | 5 +++--
 tests/test_jwt_identity.py       | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/aiohttp_security/jwt_identity.py b/aiohttp_security/jwt_identity.py
index bd6ffe7..d3481b4 100644
--- a/aiohttp_security/jwt_identity.py
+++ b/aiohttp_security/jwt_identity.py
@@ -15,11 +15,12 @@ AUTH_SCHEME = 'Bearer '
 
 
 class JWTIdentityPolicy(AbstractIdentityPolicy):
-    def __init__(self, secret, algorithm='HS256'):
+    def __init__(self, secret, algorithm='HS256', key: str = 'login'):
         if jwt is None:
             raise RuntimeError('Please install `PyJWT`')
         self.secret = secret
         self.algorithm = algorithm
+        self.key = key
 
     async def identify(self, request):
         header_identity = request.headers.get(AUTH_HEADER_NAME)
@@ -36,7 +37,7 @@ class JWTIdentityPolicy(AbstractIdentityPolicy):
         identity = jwt.decode(token,
                               self.secret,
                               algorithms=[self.algorithm])
-        return identity
+        return identity.get(self.key)
 
     async def remember(self, *args, **kwargs):  # pragma: no cover
         pass
diff --git a/tests/test_jwt_identity.py b/tests/test_jwt_identity.py
index 813deac..83ae9d6 100644
--- a/tests/test_jwt_identity.py
+++ b/tests/test_jwt_identity.py
@@ -43,7 +43,7 @@ async def test_identify(loop, make_token, aiohttp_client):
     async def check(request):
         policy = request.app[IDENTITY_KEY]
         identity = await policy.identify(request)
-        assert 'Andrew' == identity['login']
+        assert 'Andrew' == identity
         return web.Response()
 
     app = web.Application()