fix: fix exception handling bug in tasklets (#520)

Fixes #519

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/tasklets.py

@@ -309,10 +309,14 @@ def _advance_tasklet(self, send_value=None, error=None):
                         traceback = error.__traceback__
                     except AttributeError:  # pragma: NO PY3 COVER  # pragma: NO BRANCH  # noqa: E501
                         traceback = None
-                    self.generator.throw(type(error), error, traceback)
 
-                # send_value will be None if this is the first time
-                yielded = self.generator.send(send_value)
+                    yielded = self.generator.throw(
+                        type(error), error, traceback
+                    )
+
+                else:
+                    # send_value will be None if this is the first time
+                    yielded = self.generator.send(send_value)
 
                 # Context may have changed in tasklet
                 self.context = context_module.get_context()

packages/google-cloud-ndb/noxfile.py

@@ -28,6 +28,7 @@
 ALL_INTERPRETERS = ("2.7", "3.6", "3.7")
 PY3_INTERPRETERS = ("3.6", "3.7")
 MAJOR_INTERPRETERS = ("2.7", "3.7")
+BLACK_VERSION = "black==19.10b0"
 
 
 def get_path(*names):
@@ -100,15 +101,15 @@ def lint(session):
     Returns a failure if the linters find linting errors or sufficiently
     serious code quality issues.
     """
-    session.install("flake8", "black")
+    session.install("flake8", BLACK_VERSION)
     run_black(session, use_check=True)
     session.run("flake8", "google", "tests")
 
 
 @nox.session(py=DEFAULT_INTERPRETER)
 def blacken(session):
     # Install all dependencies.
-    session.install("black")
+    session.install(BLACK_VERSION)
     # Run ``black``.
     run_black(session)
 

packages/google-cloud-ndb/tests/unit/test__retry.py

@@ -25,6 +25,8 @@
 from google.cloud.ndb import _retry
 from google.cloud.ndb import tasklets
 
+from . import utils
+
 
 class Test_retry:
     @staticmethod
@@ -79,6 +81,40 @@ def test_transient_error(core_retry, sleep):
 
         sleep.assert_called_once_with(0)
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb.tasklets.sleep")
+    @mock.patch("google.cloud.ndb._retry.core_retry")
+    def test_transient_error_callback_is_tasklet(core_retry, sleep):
+        """Regression test for #519
+
+        https://github.com/googleapis/python-ndb/issues/519
+        """
+        core_retry.exponential_sleep_generator.return_value = itertools.count()
+        core_retry.if_transient_error.return_value = True
+
+        sleep_future = tasklets.Future("sleep")
+        sleep.return_value = sleep_future
+
+        callback = mock.Mock(
+            side_effect=[
+                utils.future_exception(Exception("Spurious error.")),
+                utils.future_result("foo"),
+            ]
+        )
+        retry = _retry.retry_async(callback)
+        future = retry()
+
+        # This is the important check for the bug in #519. We need to make sure
+        # that we're waiting for the sleep future to complete before moving on.
+        assert future.running()
+
+        # Finish sleeping
+        sleep_future.set_result(None)
+        assert future.result() == "foo"
+
+        sleep.assert_called_once_with(0)
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     @mock.patch("google.cloud.ndb.tasklets.sleep")

packages/google-cloud-ndb/tests/unit/test_tasklets.py

@@ -371,6 +371,26 @@ def generator_function(dependency):
         with pytest.raises(Exception):
             future.result()
 
+    @staticmethod
+    def test__advance_tasklet_dependency_raises_with_try_except(in_context):
+        def generator_function(dependency, error_handler):
+            try:
+                yield dependency
+            except Exception:
+                result = yield error_handler
+                raise tasklets.Return(result)
+
+        error = Exception("Spurious error.")
+        dependency = tasklets.Future()
+        error_handler = tasklets.Future()
+        generator = generator_function(dependency, error_handler)
+        future = tasklets._TaskletFuture(generator, in_context)
+        future._advance_tasklet()
+        dependency.set_exception(error)
+        assert future.running()
+        error_handler.set_result("hi mom!")
+        assert future.result() == "hi mom!"
+
     @staticmethod
     def test__advance_tasklet_yields_rpc(in_context):
         def generator_function(dependency):

packages/google-cloud-ndb/tests/unit/utils.py

@@ -35,6 +35,13 @@ def future_result(result):
     return future
 
 
+def future_exception(exception):
+    """Return a future with the given result."""
+    future = tasklets.Future()
+    future.set_exception(exception)
+    return future
+
+
 def future_results(*results):
     """Return a sequence of futures for the given results."""
     return [future_result(result) for result in results]