def testInitNewUser(self):
users.get_current_user = mock.MagicMock(return_value=self.new_app_user)
self.assertEqual(self.new_app_user, InitUser('new token').user)
new_app_user_ndb = AppUser.query(
AppUser.user == users.get_current_user()).fetch()[0]
self.assertTrue(new_app_user_ndb)
self.assertEqual('new token', new_app_user_ndb.refresh_token)
def InitUser(refresh_token=None):
"""Initialize application user.
Retrieve existing user credentials from datastore or add new user.
If a refresh token is specified, it replaces the old refresh token and the
date_acquired is updated as well.
Args:
refresh_token: str A new refresh token received from the auth flow.
Defaults to None.
Returns:
AppUser instance of the application user.
"""
result = AppUser.query(
AppUser.user == users.get_current_user()).fetch(limit=1)
if result:
# app_user exists
app_user = result[0]
if not app_user.refresh_token:
# update refresh_token if provided in the arguments
app_user.refresh_token = refresh_token
app_user.date_acquired = datetime.datetime.now()
else:
app_user = AppUser(
user=users.get_current_user(),
email=users.get_current_user().email(),
refresh_token=refresh_token)
app_user.put()
return app_user
def RevokeOldCredentials():
"""Revoke old credentials.
Find all users in the datastore with refresh tokens older than 30 days and
revokes them.
"""
users_with_expired_tokens = AppUser.query(
AppUser.date_acquired <=
datetime.datetime.now() - datetime.timedelta(30)).fetch()
for expired_user in users_with_expired_tokens:
expired_user.refresh_token = None
ndb.put_multi(users_with_expired_tokens)
def get_user_key(user):
return str(AppUser.query(AppUser.identity == user.user_id()).get().key.id())
