Implement get_event_loop. (googleapis#6456)

Chris Rossi · web-flow · commit 34f97881096b · 2018-11-08T17:16:58.000-05:00
Implement `get_event_loop`.
diff --git a/ndb/src/google/cloud/ndb/_eventloop.py b/ndb/src/google/cloud/ndb/_eventloop.py
@@ -21,6 +21,8 @@
 import threading
 import time
 
+from google.cloud.ndb import exceptions
+
 __all__ = [
     "add_idle",
     "async_context",
@@ -304,8 +306,9 @@ def async_context():
     context. Upon exiting the context, execution will block until all
     asynchronous calls loaded onto the event loop have finished execution.
 
-    Code within an asynchronous context should be single threaded. Internally, a
-    :class:`threading.local` instance is used to track the current event loop.
+    Code within an asynchronous context should be single threaded. Internally,
+    a :class:`threading.local` instance is used to track the current event
+    loop.
 
     In the context of a web application, it is recommended that a single
     asynchronous context be used per HTTP request. This can typically be
@@ -321,11 +324,28 @@ def async_context():
     contexts.pop()
 
 
-def add_idle(*args, **kwargs):
-    raise NotImplementedError
+def get_event_loop():
+    """Get the current event loop.
 
+    This function should be called within a context established by
+    :func:`google.cloud.ndb.async_context`.
 
-def get_event_loop(*args, **kwargs):
+    Returns:
+        EventLoop: The event loop for the current
+            context.
+
+    Raises:
+        exceptions.AsyncContextError: If called outside of a context
+            established by :func:`google.cloud.ndb.async_context`.
+    """
+    loop = contexts.current()
+    if loop:
+        return loop
+
+    raise exceptions.AsyncContextError()
+
+
+def add_idle(*args, **kwargs):
     raise NotImplementedError
 
 
diff --git a/ndb/src/google/cloud/ndb/exceptions.py b/ndb/src/google/cloud/ndb/exceptions.py
@@ -22,6 +22,7 @@
 
 __all__ = [
     "Error",
+    "AsyncContextError",
     "BadValueError",
     "BadArgumentError",
     "Rollback",
@@ -33,6 +34,20 @@ class Error(Exception):
     """Base datastore error type."""
 
 
+class AsyncContextError(Error):
+    """Indicates an async call being made without a context.
+
+    Raised whenever an asynchronous call is made outside of a context
+    established by :func:`google.cloud.ndb.async_context`.
+    """
+
+    def __init__(self):
+        super(AsyncContextError, self).__init__(
+            "No currently running event loop. Asynchronous calls must be made "
+            "in context established by google.cloud.ndb.async_context."
+        )
+
+
 class BadValueError(Error):
     """Indicates a property value or filter value is invalid.
 
diff --git a/ndb/tests/unit/test__eventloop.py b/ndb/tests/unit/test__eventloop.py
@@ -17,9 +17,11 @@
 
 import pytest
 
-from google.cloud.ndb import _eventloop as eventloop
 import tests.unit.utils
 
+from google.cloud.ndb import exceptions
+from google.cloud.ndb import _eventloop as eventloop
+
 
 def test___all__():
     tests.unit.utils.verify___all__(eventloop)
@@ -317,14 +319,16 @@ def test_async_context(EventLoop):
     one.run.assert_called_once_with()
 
 
-def test_add_idle():
-    with pytest.raises(NotImplementedError):
-        eventloop.add_idle()
+def test_get_event_loop():
+    with pytest.raises(exceptions.AsyncContextError):
+        eventloop.get_event_loop()
+    with eventloop.async_context():
+        assert isinstance(eventloop.get_event_loop(), eventloop.EventLoop)
 
 
-def test_get_event_loop():
+def test_add_idle():
     with pytest.raises(NotImplementedError):
-        eventloop.get_event_loop()
+        eventloop.add_idle()
 
 
 def test_queue_call():
