Adjust device flow exit_condition signature and behavior

Author: rayluo

oauth2cli/oauth2.py

@@ -197,6 +197,7 @@ def initiate_device_flow(self, scope=None, timeout=None, **kwargs):
             **kwargs).json()
         flow["interval"] = int(flow.get("interval", 5))  # Some IdP returns string
         flow["expires_in"] = int(flow.get("expires_in", 1800))
+        flow["expires_at"] = time.time() + flow["expires_in"]  # We invent this
         return flow
 
     def _obtain_token_by_device_flow(self, flow, **kwargs):
@@ -219,25 +220,41 @@ def _obtain_token_by_device_flow(self, flow, **kwargs):
         flow["latest_attempt_at"] = now
         return result
 
-    def obtain_token_by_device_flow(self, flow, exit_condition=lambda: True, **kwargs):
+    def obtain_token_by_device_flow(self,
+            flow,
+            exit_condition=lambda flow: flow.get("expires_at", 0) < time.time(),
+            **kwargs):
         # type: (dict, Callable) -> dict
-        """Obtain token by a device flow object, with optional polling effect.
+        """Obtain token by a device flow object, with customizable polling effect.
 
         Args:
             flow (dict):
                 An object previously generated by initiate_device_flow(...).
+                Its content WILL BE CHANGED by this method during each run.
+                We share this object with you, so that you could implement
+                your own loop, should you choose to do so.
+
             exit_condition (Callable):
                 This method implements a loop to provide polling effect.
                 The loop's exit condition is calculated by this callback.
-                The default callback makes the loop run only once, i.e. no polling.
+
+                The default callback makes the loop run until the flow expires.
+                Therefore, one of the ways to exit the polling early,
+                is to change the flow["expires_at"] to a small number such as 0.
+
+                In case you are doing async programming, you may want to
+                completely turn off the loop. You can do so by using a callback as:
+
+                    exit_condition = lambda flow: True
+
+                to make the loop run only once, i.e. no polling, hence non-block.
         """
-        _flow = flow.copy()
         while True:
-            result = self._obtain_token_by_device_flow(_flow, **kwargs)
+            result = self._obtain_token_by_device_flow(flow, **kwargs)
             if result.get("error") not in self.DEVICE_FLOW_RETRIABLE_ERRORS:
                 return result
-            for i in range(_flow.get("interval", 5)):  # Wait interval seconds
-                if exit_condition():
+            for i in range(flow.get("interval", 5)):  # Wait interval seconds
+                if exit_condition(flow):
                     return result
                 time.sleep(1)  # Shorten each round, to make exit more responsive
 

tests/test_client.py

@@ -147,10 +147,8 @@ def test_device_flow(self):
 
         duration = 30
         logger.warn("We will wait up to %d seconds for you to sign in" % duration)
-        result = self.client.obtain_token_by_device_flow(
-            flow,
-            exit_condition=lambda end=time.time() + duration: time.time() > end)
-
+        flow["expires_at"] = time.time() + duration  # Shorten the time for quick test
+        result = self.client.obtain_token_by_device_flow(flow)
         self.assertLoosely(
                 result,
                 assertion=lambda: self.assertIn('access_token', result),