text: # Copyright 2018 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  try:
      from unittest import mock
  except ImportError:  # pragma: NO PY3 COVER
      import mock
  
  import pytest
  
  from google.cloud.ndb import context as context_module
  from google.cloud.ndb import _eventloop
  from google.cloud.ndb import exceptions
  from google.cloud.ndb import _remote
  from google.cloud.ndb import tasklets
  
  from . import utils
  
  
  def test___all__():
      utils.verify___all__(tasklets)
  
  
  def test_add_flow_exception():
      with pytest.raises(NotImplementedError):
          tasklets.add_flow_exception()
  
  
  class TestFuture:
      @staticmethod
      def test_constructor():
          future = tasklets.Future()
          assert future.running()
          assert not future.done()
          assert future.info == "Unknown"
  
      @staticmethod
      def test_constructor_w_info():
          future = tasklets.Future("Testing")
          assert future.running()
          assert not future.done()
          assert future.info == "Testing"
  
      @staticmethod
      def test___repr__():
          future = tasklets.Future("The Children")
          assert repr(future) == "Future('The Children') <{}>".format(id(future))
  
      @staticmethod
      def test_set_result():
          future = tasklets.Future()
          future.set_result(42)
          assert future.result() == 42
          assert future.get_result() == 42
          assert future.done()
          assert not future.running()
          assert future.exception() is None
          assert future.get_exception() is None
          assert future.get_traceback() is None
  
      @staticmethod
      def test_set_result_already_done():
          future = tasklets.Future()
          future.set_result(42)
          with pytest.raises(RuntimeError):
              future.set_result(42)
  
      @staticmethod
      def test_add_done_callback():
          callback1 = mock.Mock()
          callback2 = mock.Mock()
          future = tasklets.Future()
          future.add_done_callback(callback1)
          future.add_done_callback(callback2)
          future.set_result(42)
  
          callback1.assert_called_once_with(future)
          callback2.assert_called_once_with(future)
  
      @staticmethod
      def test_add_done_callback_already_done():
          callback = mock.Mock()
          future = tasklets.Future()
          future.set_result(42)
          future.add_done_callback(callback)
          callback.assert_called_once_with(future)
  
      @staticmethod
      def test_set_exception():
          future = tasklets.Future()
          error = Exception("Spurious Error")
          future.set_exception(error)
          assert future.exception() is error
          assert future.get_exception() is error
          assert future.get_traceback() is getattr(error, "__traceback__", None)
          with pytest.raises(Exception):
              future.result()
  
      @staticmethod
      def test_set_exception_with_callback():
          callback = mock.Mock()
          future = tasklets.Future()
          future.add_done_callback(callback)
          error = Exception("Spurious Error")
          future.set_exception(error)
          assert future.exception() is error
          assert future.get_exception() is error
          assert future.get_traceback() is getattr(error, "__traceback__", None)
          callback.assert_called_once_with(future)
  
      @staticmethod
      def test_set_exception_already_done():
          future = tasklets.Future()
          error = Exception("Spurious Error")
          future.set_exception(error)
          with pytest.raises(RuntimeError):
              future.set_exception(error)
  
      @staticmethod
      @mock.patch("google.cloud.ndb.tasklets._eventloop")
      def test_wait(_eventloop):
          def side_effects(future):
              yield True
              yield True
              future.set_result(42)
              yield True
  
          future = tasklets.Future()
          _eventloop.run1.side_effect = side_effects(future)
          future.wait()
          assert future.result() == 42
          assert _eventloop.run1.call_count == 3
  
      @staticmethod
      @mock.patch("google.cloud.ndb.tasklets._eventloop")
      def test_wait_loop_exhausted(_eventloop):
          future = tasklets.Future()
          _eventloop.run1.return_value = False
          with pytest.raises(RuntimeError):
              future.wait()
  
      @staticmethod
      @mock.patch("google.cloud.ndb.tasklets._eventloop")
      def test_check_success(_eventloop):
          def side_effects(future):
              yield True
              yield True
              future.set_result(42)
              yield True
  
          future = tasklets.Future()
          _eventloop.run1.side_effect = side_effects(future)
          future.check_success()
          assert future.result() == 42
          assert _eventloop.run1.call_count == 3
  
      @staticmethod
      @mock.patch("google.cloud.ndb.tasklets._eventloop")
      def test_check_success_failure(_eventloop):
          error = Exception("Spurious error")
  
          def side_effects(future):
              yield True
              yield True
              future.set_exception(error)
              yield True
  
          future = tasklets.Future()
          _eventloop.run1.side_effect = side_effects(future)
          with pytest.raises(Exception) as error_context:
              future.check_success()
  
          assert error_context.value is error
  
      @staticmethod
      @mock.patch("google.cloud.ndb.tasklets._eventloop")
      def test_result_block_for_result(_eventloop):
          def side_effects(future):
              yield True
              yield True
              future.set_result(42)
              yield True
  
          future = tasklets.Future()
          _eventloop.run1.side_effect = side_effects(future)
          assert future.result() == 42
          assert _eventloop.run1.call_count == 3
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_cancel():
          # Integration test. Actually test that a cancel propagates properly.
          rpc = tasklets.Future("Fake RPC")
          wrapped_rpc = _remote.RemoteCall(rpc, "Wrapped Fake RPC")
  
          @tasklets.tasklet
          def inner_tasklet():
              yield wrapped_rpc
  
          @tasklets.tasklet
          def outer_tasklet():
              yield inner_tasklet()
  
          future = outer_tasklet()
          assert not future.cancelled()
          future.cancel()
          assert rpc.cancelled()
  
          with pytest.raises(exceptions.Cancelled):
              future.result()
  
          assert future.cancelled()
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_cancel_already_done():
          future = tasklets.Future("testing")
          future.set_result(42)
          future.cancel()  # noop
          assert not future.cancelled()
          assert future.result() == 42
  
      @staticmethod
      def test_cancelled():
          future = tasklets.Future()
          assert future.cancelled() is False
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_wait_any():
          futures = [tasklets.Future() for _ in range(3)]
  
          def callback():
              futures[1].set_result(42)
  
          _eventloop.add_idle(callback)
  
          future = tasklets.Future.wait_any(futures)
          assert future is futures[1]
          assert future.result() == 42
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_wait_any_loop_exhausted():
          futures = [tasklets.Future() for _ in range(3)]
  
          with pytest.raises(RuntimeError):
              tasklets.Future.wait_any(futures)
  
      @staticmethod
      def test_wait_any_no_futures():
          assert tasklets.Future.wait_any(()) is None
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_wait_all():
          futures = [tasklets.Future() for _ in range(3)]
  
          def make_callback(index, result):
              def callback():
                  futures[index].set_result(result)
  
              return callback
  
          _eventloop.add_idle(make_callback(0, 42))
          _eventloop.add_idle(make_callback(1, 43))
          _eventloop.add_idle(make_callback(2, 44))
  
          tasklets.Future.wait_all(futures)
          assert futures[0].done()
          assert futures[0].result() == 42
          assert futures[1].done()
          assert futures[1].result() == 43
          assert futures[2].done()
          assert futures[2].result() == 44
  
      @staticmethod
      def test_wait_all_no_futures():
          assert tasklets.Future.wait_all(()) is None
  
  
  class Test_TaskletFuture:
      @staticmethod
      def test_constructor():
          generator = object()
          context = object()
          future = tasklets._TaskletFuture(generator, context)
          assert future.generator is generator
          assert future.context is context
          assert future.info == "Unknown"
  
      @staticmethod
      def test___repr__():
          future = tasklets._TaskletFuture(None, None, info="Female")
          assert repr(future) == "_TaskletFuture('Female') <{}>".format(id(future))
  
      @staticmethod
      def test__advance_tasklet_return(in_context):
          def generator_function():
              yield
              raise tasklets.Return(42)
  
          generator = generator_function()
          next(generator)  # skip ahead to return
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          assert future.result() == 42
  
      @staticmethod
      def test__advance_tasklet_generator_raises(in_context):
          error = Exception("Spurious error.")
  
          def generator_function():
              yield
              raise error
  
          generator = generator_function()
          next(generator)  # skip ahead to return
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          assert future.exception() is error
  
      @staticmethod
      def test__advance_tasklet_bad_yield(in_context):
          def generator_function():
              yield 42
  
          generator = generator_function()
          future = tasklets._TaskletFuture(generator, in_context)
          with pytest.raises(RuntimeError):
              future._advance_tasklet()
  
      @staticmethod
      def test__advance_tasklet_dependency_returns(in_context):
          def generator_function(dependency):
              some_value = yield dependency
              raise tasklets.Return(some_value + 42)
  
          dependency = tasklets.Future()
          generator = generator_function(dependency)
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          dependency.set_result(21)
          assert future.result() == 63
  
      @staticmethod
      def test__advance_tasklet_dependency_raises(in_context):
          def generator_function(dependency):
              yield dependency
  
          error = Exception("Spurious error.")
          dependency = tasklets.Future()
          generator = generator_function(dependency)
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          dependency.set_exception(error)
          assert future.exception() is error
          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):
              value = yield dependency
              raise tasklets.Return(value + 3)
  
          dependency = mock.Mock(spec=_remote.RemoteCall)
          dependency.exception.return_value = None
          dependency.result.return_value = 8
          generator = generator_function(dependency)
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
  
          callback = dependency.add_done_callback.call_args[0][0]
          callback(dependency)
          _eventloop.run()
          assert future.result() == 11
  
      @staticmethod
      def test__advance_tasklet_parallel_yield(in_context):
          def generator_function(dependencies):
              one, two = yield dependencies
              raise tasklets.Return(one + two)
  
          dependencies = (tasklets.Future(), tasklets.Future())
          generator = generator_function(dependencies)
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          dependencies[0].set_result(8)
          dependencies[1].set_result(3)
          assert future.result() == 11
          assert future.context is in_context
  
      @staticmethod
      def test_cancel_not_waiting(in_context):
          dependency = tasklets.Future()
          future = tasklets._TaskletFuture(None, in_context)
          future.cancel()
  
          assert not dependency.cancelled()
          with pytest.raises(exceptions.Cancelled):
              future.result()
  
      @staticmethod
      def test_cancel_waiting_on_dependency(in_context):
          def generator_function(dependency):
              yield dependency
  
          dependency = tasklets.Future()
          generator = generator_function(dependency)
          future = tasklets._TaskletFuture(generator, in_context)
          future._advance_tasklet()
          future.cancel()
  
          assert dependency.cancelled()
          with pytest.raises(exceptions.Cancelled):
              future.result()
  
  
  class Test_MultiFuture:
      @staticmethod
      def test___repr__():
          this, that = (tasklets.Future("this"), tasklets.Future("that"))
          future = tasklets._MultiFuture((this, that))
          assert repr(future) == (
              "_MultiFuture(Future('this') <{}>,"
              " Future('that') <{}>) <{}>".format(id(this), id(that), id(future))
          )
  
      @staticmethod
      def test_success():
          dependencies = (tasklets.Future(), tasklets.Future())
          future = tasklets._MultiFuture(dependencies)
          dependencies[0].set_result("one")
          dependencies[1].set_result("two")
          assert future.result() == ("one", "two")
  
      @staticmethod
      def test_error():
          dependencies = (tasklets.Future(), tasklets.Future())
          future = tasklets._MultiFuture(dependencies)
          error = Exception("Spurious error.")
          dependencies[0].set_exception(error)
          dependencies[1].set_result("two")
          assert future.exception() is error
          with pytest.raises(Exception):
              future.result()
  
      @staticmethod
      def test_cancel():
          dependencies = (tasklets.Future(), tasklets.Future())
          future = tasklets._MultiFuture(dependencies)
          future.cancel()
          assert dependencies[0].cancelled()
          assert dependencies[1].cancelled()
          with pytest.raises(exceptions.Cancelled):
              future.result()
  
      @staticmethod
      def test_no_dependencies():
          future = tasklets._MultiFuture(())
          assert future.result() == ()
  
      @staticmethod
      def test_nested():
          dependencies = [tasklets.Future() for _ in range(3)]
          future = tasklets._MultiFuture((dependencies[0], dependencies[1:]))
          for i, dependency in enumerate(dependencies):
              dependency.set_result(i)
  
          assert future.result() == (0, (1, 2))
  
  
  class Test__get_return_value:
      @staticmethod
      def test_no_args():
          stop = StopIteration()
          assert tasklets._get_return_value(stop) is None
  
      @staticmethod
      def test_one_arg():
          stop = StopIteration(42)
          assert tasklets._get_return_value(stop) == 42
  
      @staticmethod
      def test_two_args():
          stop = StopIteration(42, 21)
          assert tasklets._get_return_value(stop) == (42, 21)
  
  
  class Test_tasklet:
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_generator():
          @tasklets.tasklet
          def generator(dependency):
              value = yield dependency
              raise tasklets.Return(value + 3)
  
          dependency = tasklets.Future()
          future = generator(dependency)
          assert isinstance(future, tasklets._TaskletFuture)
          dependency.set_result(8)
          assert future.result() == 11
  
      # Can't make this work with 2.7, because the return with argument inside
      # generator error crashes the pytest collection process, even with skip
      # @staticmethod
      # @pytest.mark.skipif(sys.version_info[0] == 2, reason="requires python3")
      # @pytest.mark.usefixtures("in_context")
      # def test_generator_using_return():
      #     @tasklets.tasklet
      #     def generator(dependency):
      #         value = yield dependency
      #         return value + 3
  
      #     dependency = tasklets.Future()
      #     future = generator(dependency)
      #     assert isinstance(future, tasklets._TaskletFuture)
      #     dependency.set_result(8)
      #     assert future.result() == 11
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_regular_function():
          @tasklets.tasklet
          def regular_function(value):
              return value + 3
  
          future = regular_function(8)
          assert isinstance(future, tasklets.Future)
          assert future.result() == 11
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_regular_function_raises_Return():
          @tasklets.tasklet
          def regular_function(value):
              raise tasklets.Return(value + 3)
  
          future = regular_function(8)
          assert isinstance(future, tasklets.Future)
          assert future.result() == 11
  
      @staticmethod
      def test_context_management(in_context):
          @tasklets.tasklet
          def some_task(transaction, future):
              assert context_module.get_context().transaction == transaction
              yield future
              raise tasklets.Return(context_module.get_context().transaction)
  
          future_foo = tasklets.Future("foo")
          with in_context.new(transaction="foo").use():
              task_foo = some_task("foo", future_foo)
  
          future_bar = tasklets.Future("bar")
          with in_context.new(transaction="bar").use():
              task_bar = some_task("bar", future_bar)
  
          future_foo.set_result(None)
          future_bar.set_result(None)
  
          assert task_foo.result() == "foo"
          assert task_bar.result() == "bar"
  
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_context_changed_in_tasklet():
          @tasklets.tasklet
          def some_task(transaction, future1, future2):
              context = context_module.get_context()
              assert context.transaction is None
              with context.new(transaction=transaction).use():
                  assert context_module.get_context().transaction == transaction
                  yield future1
                  assert context_module.get_context().transaction == transaction
                  yield future2
                  assert context_module.get_context().transaction == transaction
              assert context_module.get_context() is context
  
          future_foo1 = tasklets.Future("foo1")
          future_foo2 = tasklets.Future("foo2")
          task_foo = some_task("foo", future_foo1, future_foo2)
  
          future_bar1 = tasklets.Future("bar1")
          future_bar2 = tasklets.Future("bar2")
          task_bar = some_task("bar", future_bar1, future_bar2)
  
          future_foo1.set_result(None)
          future_bar1.set_result(None)
          future_foo2.set_result(None)
          future_bar2.set_result(None)
  
          task_foo.check_success()
          task_bar.check_success()
  
  
  class Test_wait_any:
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_it():
          futures = [tasklets.Future() for _ in range(3)]
  
          def callback():
              futures[1].set_result(42)
  
          _eventloop.add_idle(callback)
  
          future = tasklets.wait_any(futures)
          assert future is futures[1]
          assert future.result() == 42
  
      @staticmethod
      def test_it_no_futures():
          assert tasklets.wait_any(()) is None
  
  
  class Test_wait_all:
      @staticmethod
      @pytest.mark.usefixtures("in_context")
      def test_it():
          futures = [tasklets.Future() for _ in range(3)]
  
          def make_callback(index, result):
              def callback():
                  futures[index].set_result(result)
  
              return callback
  
          _eventloop.add_idle(make_callback(0, 42))
          _eventloop.add_idle(make_callback(1, 43))
          _eventloop.add_idle(make_callback(2, 44))
  
          tasklets.wait_all(futures)
          assert futures[0].done()
          assert futures[0].result() == 42
          assert futures[1].done()
          assert futures[1].result() == 43
          assert futures[2].done()
          assert futures[2].result() == 44
  
      @staticmethod
      def test_it_no_futures():
          assert tasklets.wait_all(()) is None
  
  
  @pytest.mark.usefixtures("in_context")
  @mock.patch("google.cloud.ndb._eventloop.time")
  def test_sleep(time_module, context):
      time_module.time.side_effect = [0, 0, 1]
      future = tasklets.sleep(1)
      assert future.get_result() is None
      time_module.sleep.assert_called_once_with(1)
  
  
  def test_make_context():
      with pytest.raises(NotImplementedError):
          tasklets.make_context()
  
  
  def test_make_default_context():
      with pytest.raises(NotImplementedError):
          tasklets.make_default_context()
  
  
  class TestQueueFuture:
      @staticmethod
      def test_constructor():
          with pytest.raises(NotImplementedError):
              tasklets.QueueFuture()
  
  
  class TestReducingFuture:
      @staticmethod
      def test_constructor():
          with pytest.raises(NotImplementedError):
              tasklets.ReducingFuture()
  
  
  def test_Return():
      assert not issubclass(tasklets.Return, StopIteration)
      assert issubclass(tasklets.Return, Exception)
  
  
  class TestSerialQueueFuture:
      @staticmethod
      def test_constructor():
          with pytest.raises(NotImplementedError):
              tasklets.SerialQueueFuture()
  
  
  def test_set_context():
      with pytest.raises(NotImplementedError):
          tasklets.set_context()
  
  
  @pytest.mark.usefixtures("in_context")
  def test_synctasklet():
      @tasklets.synctasklet
      def generator_function(value):
          future = tasklets.Future(value)
          future.set_result(value)
          x = yield future
          raise tasklets.Return(x + 3)
  
      result = generator_function(8)
      assert result == 11
  
  
  @pytest.mark.usefixtures("in_context")
  def test_toplevel():
      @tasklets.toplevel
      def generator_function(value):
          future = tasklets.Future(value)
          future.set_result(value)
          x = yield future
          raise tasklets.Return(x + 3)
  
      idle = mock.Mock(__name__="idle", return_value=None)
      _eventloop.add_idle(idle)
  
      result = generator_function(8)
      assert result == 11
      idle.assert_called_once_with()
  