Refresh token with less that 30s to live (#134)

* Refresh token with less that 30s to live

* Fix tests

neptune/oauth.py

@@ -27,12 +27,13 @@ class NeptuneAuth(AuthBase):
 
     def __init__(self, session):
         self.session = session
+        self.token_expires_at = 0
 
     def __call__(self, r):
         try:
             return self._add_token(r)
         except TokenExpiredError:
-            self.session.refresh_token(self.session.auto_refresh_url)
+            self._refresh_token()
             return self._add_token(r)
 
     def _add_token(self, r):
@@ -43,6 +44,16 @@ def _add_token(self, r):
                                                                   headers=r.headers)
         return r
 
+    def refresh_token_if_needed(self):
+        if self.token_expires_at - time.time() < 30:
+            self._refresh_token()
+
+    def _refresh_token(self):
+        self.session.refresh_token(self.session.auto_refresh_url)
+        if self.session.token is not None and self.session.token.get('access_token') is not None:
+            decoded_json_token = jwt.decode(self.session.token.get('access_token'), verify=False)
+            self.token_expires_at = decoded_json_token.get(u'exp')
+
 
 class NeptuneAuthenticator(Authenticator):
 
@@ -69,6 +80,7 @@ def matches(self, url):
         return True
 
     def apply(self, request):
+        self.auth.refresh_token_if_needed()
         request.auth = self.auth
         return request
 

tests/neptune/test_oauth.py

@@ -31,7 +31,9 @@ def setUp(self):
         super(TestNeptuneAuth, self).setUp()
 
         self.session = MagicMock()
+        self.session.token = dict()
         self.neptune_auth = NeptuneAuth(self.session)
+        self.neptune_auth.token_expires_at = time.time() + 60
         self.request = a_request()
 
         self.url, self.method, self.body, self.headers = \
@@ -97,6 +99,7 @@ def test_apply_oauth2_session_to_request(self, time_mock, session_mock):
         # and
         session = MagicMock()
         session_mock.return_value = session
+        session.token = dict()
 
         # and
         neptune_authenticator = NeptuneAuthenticator(auth_tokens)