Allow retry timer callback to signal complete

tests/test_retry.py

@@ -17,6 +17,7 @@ class State(AbstractContextManager):
         self.callbacks = 0
         self.timed_out = False
         self.tag = str(uuid.uuid4())
+        self.results = [self.tag, self.tag, self.tag]
         self.got_tag = None
         assert self.retry_thread.is_alive()
 
@@ -30,7 +31,7 @@ class State(AbstractContextManager):
         self.tries = tries
         self.got_tag = tag
         self.callbacks += 1
-        return self.tag
+        return self.results[tries - 1]
 
     def timeout(self, tag, tries):
         self.tries = tries
@@ -47,11 +48,11 @@ class State(AbstractContextManager):
 def test_retry_timeout():
 
     with State(0.1) as state:
-        state.retry_thread.begin(try_limit=3,
-                                wait_delay=state.delay,
-                                try_callback=state.retry,
-                                timeout_callback=state.timeout)
-
+        return_tag = state.retry_thread.begin(try_limit=3,
+                                              wait_delay=state.delay,
+                                              try_callback=state.retry,
+                                              timeout_callback=state.timeout)
+        assert return_tag == state.tag
         assert state.tries == 1
         assert state.callbacks == 1
         assert state.got_tag is None
@@ -81,15 +82,62 @@ def test_retry_timeout():
         assert state.tries == 3
 
 
-def test_retry_complete():
+def test_retry_immediate_complete():
     with State(0.01) as state:
-        state.retry_thread.begin(try_limit=3,
-                                 wait_delay=state.delay,
-                                 try_callback=state.retry,
-                                 timeout_callback=state.timeout)
-
+        state.results[0] = False
+        return_tag = state.retry_thread.begin(try_limit=3,
+                                              wait_delay=state.delay,
+                                              try_callback=state.retry,
+                                              timeout_callback=state.timeout)
+        assert not return_tag
+        assert state.callbacks == 1
+        assert not state.got_tag
+        assert not state.timed_out
+        time.sleep(state.delay * 3)
         assert state.tries == 1
         assert state.callbacks == 1
+        assert not state.got_tag
+        assert not state.timed_out
+
+
+def test_retry_return_complete():
+    with State(0.01) as state:
+        state.results[1] = False
+        return_tag = state.retry_thread.begin(try_limit=3,
+                                              wait_delay=state.delay,
+                                              try_callback=state.retry,
+                                              timeout_callback=state.timeout)
+        assert return_tag == state.tag
+        assert state.callbacks == 1
+        assert state.got_tag is None
+        assert not state.timed_out
+        time.sleep(state.delay / 2.0)
+        time.sleep(state.delay)
+        assert state.tries == 2
+        assert state.callbacks == 2
+        assert state.got_tag == state.tag
+        assert not state.timed_out
+
+        time.sleep(state.delay)
+        assert state.tries == 2
+        assert state.callbacks == 2
+        assert state.got_tag == state.tag
+        assert not state.timed_out
+
+        # check no more callbacks
+        time.sleep(state.delay * 3.0)
+        assert state.callbacks == 2
+        assert state.tries == 2
+
+
+def test_retry_set_complete():
+    with State(0.01) as state:
+        return_tag = state.retry_thread.begin(try_limit=3,
+                                              wait_delay=state.delay,
+                                              try_callback=state.retry,
+                                              timeout_callback=state.timeout)
+        assert return_tag == state.tag
+        assert state.callbacks == 1
         assert state.got_tag is None
         assert not state.timed_out
         time.sleep(state.delay / 2.0)