Python google.appengine.ext.ndb.tasklet() Examples

def test_unordered_first_concurrent_jobs(self):
    items = [10, 5, 0]

    log = []

    @ndb.tasklet
    def fn_async(x):
      log.append('%d started' % x)
      yield ndb.sleep(float(x) / 1000)
      log.append('%d finishing' % x)
      raise ndb.Return(x)

    expected = [(5, 5), (0, 0), (10, 10)]
    actual = utils.async_apply(
        items, fn_async, concurrent_jobs=2, unordered=True)
    self.assertFalse(isinstance(actual, list))
    self.assertEqual(expected, list(actual))
    self.assertEqual(log, [
        '10 started',
        '5 started',
        '5 finishing',
        '0 started',
        '0 finishing',
        '10 finishing',
    ]) 

def _run_until_rpc():
  """Eagerly evaluate tasklets until it is blocking on some RPC.

  Usually ndb eventloop el isn't run until some code calls future.get_result().

  When an async tasklet is called, the tasklet wrapper evaluates the tasklet
  code into a generator, enqueues a callback _help_tasklet_along onto
  the el.current queue, and returns a future.

  _help_tasklet_along, when called by the el, will
  get one yielded value from the generator. If the value if another future,
  set up a callback _on_future_complete to invoke _help_tasklet_along
  when the dependent future fulfills. If the value if a RPC, set up a
  callback _on_rpc_complete to invoke _help_tasklet_along when the RPC fulfills.
  Thus _help_tasklet_along drills down
  the chain of futures until some future is blocked by RPC. El runs
  all callbacks and constantly check pending RPC status.
  """
  el = eventloop.get_event_loop()
  while el.current:
    el.run0() 

def test_cache_with_tasklets(self):
    @utils.cache
    def f():
      ndb.sleep(0).wait()  # Yield thread.
      return 1

    @ndb.tasklet
    def g():
      yield ()  # Make g a generator.
      raise ndb.Return(f())

    def test():
      ndb.Future.wait_all([(g()), (g())])

    t = threading.Thread(target=test)
    t.daemon = True
    t.start()
    t.join(1)
    if t.is_alive():
      self.fail('deadlock') 

def _check_etag(self, etag):
    """Check if etag is the same across requests to GCS.

    If self._etag is None, set it. If etag is set, check that the new
    etag equals the old one.

    In the __init__ method, we fire one HEAD and one GET request using
    ndb tasklet. One of them would return first and set the first value.

    Args:
      etag: etag from a GCS HTTP response. None if etag is not part of the
        response header. It could be None for example in the case of GCS
        composite file.

    Raises:
      ValueError: if two etags are not equal.
    """
    if etag is None:
      return
    elif self._etag is None:
      self._etag = etag
    elif self._etag != etag:
      raise ValueError('File on GCS has changed while reading.') 

def _check_etag(self, etag):
    """Check if etag is the same across requests to GCS.

    If self._etag is None, set it. If etag is set, check that the new
    etag equals the old one.

    In the __init__ method, we fire one HEAD and one GET request using
    ndb tasklet. One of them would return first and set the first value.

    Args:
      etag: etag from a GCS HTTP response. None if etag is not part of the
        response header. It could be None for example in the case of GCS
        composite file.

    Raises:
      ValueError: if two etags are not equal.
    """
    if etag is None:
      return
    elif self._etag is None:
      self._etag = etag
    elif self._etag != etag:
      raise ValueError('File on GCS has changed while reading.') 

def _run_until_rpc():
  """Eagerly evaluate tasklets until it is blocking on some RPC.

  Usually ndb eventloop el isn't run until some code calls future.get_result().

  When an async tasklet is called, the tasklet wrapper evaluates the tasklet
  code into a generator, enqueues a callback _help_tasklet_along onto
  the el.current queue, and returns a future.

  _help_tasklet_along, when called by the el, will
  get one yielded value from the generator. If the value if another future,
  set up a callback _on_future_complete to invoke _help_tasklet_along
  when the dependent future fulfills. If the value if a RPC, set up a
  callback _on_rpc_complete to invoke _help_tasklet_along when the RPC fulfills.
  Thus _help_tasklet_along drills down
  the chain of futures until some future is blocked by RPC. El runs
  all callbacks and constantly check pending RPC status.
  """
  el = eventloop.get_event_loop()
  while el.current:
    el.run0() 

def mock_methods(self):
    calls = []

    @ndb.tasklet
    def via_jwt(*args):
      calls.append(('jwt_based', args))
      raise ndb.Return({'access_token': 'token', 'exp_ts': 0})

    self.mock(service_account, '_mint_jwt_based_token_async', via_jwt)

    def via_gae_api(*args):
      calls.append(('gae_api', args))
      return 'token', 0
    self.mock(service_account.app_identity, 'get_access_token', via_gae_api)

    return calls 

def setUp(self):
    super(NetTest, self).setUp()

    @ndb.tasklet
    def get_access_token(*_args):
      raise ndb.Return(('own-token', 0))
    self.mock(auth, 'get_access_token_async', get_access_token)

    @ndb.tasklet
    def get_project_access_token(project_id, _scopes):
      raise ndb.Return(('%s-token' % project_id, 0))
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock(
        auth, 'is_internal_domain',
        lambda domain: domain == 'internal.example.com')

    self.mock(logging, 'warning', lambda *_args: None)
    self.mock(logging, 'error', lambda *_args: None) 

def test_project_tokens_fallback(self):
    @ndb.tasklet
    def get_project_access_token(_project_id, _scopes):
      raise auth.NotFoundError('not found')
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock_urlfetch([
        ({
            'deadline': 10,
            'headers': {
                # Switches to own token.
                'Authorization': 'Bearer own-token',
            },
            'method': 'GET',
            'url': 'https://external.example.com/123',
        }, Response(200, '', {})),
    ])
    net.request(url='https://external.example.com/123', project_id='project-id') 

def _check_etag(self, etag):
    """Check if etag is the same across requests to GCS.

    If self._etag is None, set it. If etag is set, check that the new
    etag equals the old one.

    In the __init__ method, we fire one HEAD and one GET request using
    ndb tasklet. One of them would return first and set the first value.

    Args:
      etag: etag from a GCS HTTP response. None if etag is not part of the
        response header. It could be None for example in the case of GCS
        composite file.

    Raises:
      ValueError: if two etags are not equal.
    """
    if etag is None:
      return
    elif self._etag is None:
      self._etag = etag
    elif self._etag != etag:
      raise ValueError('File on GCS has changed while reading.') 

def mock_memcache(self):
    local_mc = {'store': {}, 'reads': 0, 'writes': 0}

    @ndb.tasklet
    def mock_memcache_get(version, desc, part=None):
      value = local_mc['store'].get(bot_code.bot_key(version, desc, part))
      if value is not None:
        local_mc['reads'] += 1
      raise ndb.Return(value)

    @ndb.tasklet
    def mock_memcache_set(value, version, desc, part=None):
      local_mc['writes'] += 1
      key = bot_code.bot_key(version, desc, part)
      local_mc['store'][key] = value
      return ndb.Return(None)

    self.mock(bot_code, 'bot_memcache_set', mock_memcache_set)
    self.mock(bot_code, 'bot_memcache_get', mock_memcache_get)
    self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000)

    return local_mc 

def _remove_old_entity_async(key, now):
  """Removes a stale TaskDimensions or BotTaskDimensions instance.

  Returns:
    key if it was deleted.
  """
  obj = yield key.get_async()
  if not obj or obj.valid_until_ts >= now:
    raise ndb.Return(None)

  @ndb.tasklet
  def tx():
    obj = yield key.get_async()
    if obj and obj.valid_until_ts < now:
      yield key.delete_async()
      raise ndb.Return(key)

  res = yield datastore_utils.transaction_async(
      tx, propagation=ndb.TransactionOptions.INDEPENDENT)
  raise ndb.Return(res) 

def test_cache_with_tasklets(self):
    @utils.cache
    def f():
      ndb.sleep(0).wait()  # Yield thread.
      return 1

    @ndb.tasklet
    def g():
      yield ()  # Make g a generator.
      raise ndb.Return(f())

    def test():
      ndb.Future.wait_all([(g()), (g())])

    t = threading.Thread(target=test)
    t.daemon = True
    t.start()
    t.join(1)
    if t.is_alive():
      self.fail('deadlock') 

def test_unordered_first_concurrent_jobs(self):
    items = [10, 5, 0]

    log = []

    @ndb.tasklet
    def fn_async(x):
      log.append('%d started' % x)
      yield ndb.sleep(float(x) / 1000)
      log.append('%d finishing' % x)
      raise ndb.Return(x)

    expected = [(5, 5), (0, 0), (10, 10)]
    actual = utils.async_apply(
        items, fn_async, concurrent_jobs=2, unordered=True)
    self.assertFalse(isinstance(actual, list))
    self.assertEqual(expected, list(actual))
    self.assertEqual(log, [
        '10 started',
        '5 started',
        '5 finishing',
        '0 started',
        '0 finishing',
        '10 finishing',
    ]) 

def mock_methods(self):
    calls = []

    @ndb.tasklet
    def via_jwt(*args):
      calls.append(('jwt_based', args))
      raise ndb.Return({'access_token': 'token', 'exp_ts': 0})

    self.mock(service_account, '_mint_jwt_based_token_async', via_jwt)

    def via_gae_api(*args):
      calls.append(('gae_api', args))
      return 'token', 0
    self.mock(service_account.app_identity, 'get_access_token', via_gae_api)

    return calls 

def setUp(self):
    super(NetTest, self).setUp()

    @ndb.tasklet
    def get_access_token(*_args):
      raise ndb.Return(('own-token', 0))
    self.mock(auth, 'get_access_token_async', get_access_token)

    @ndb.tasklet
    def get_project_access_token(project_id, _scopes):
      raise ndb.Return(('%s-token' % project_id, 0))
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock(
        auth, 'is_internal_domain',
        lambda domain: domain == 'internal.example.com')

    self.mock(logging, 'warning', lambda *_args: None)
    self.mock(logging, 'error', lambda *_args: None) 

def test_project_tokens_fallback(self):
    @ndb.tasklet
    def get_project_access_token(_project_id, _scopes):
      raise auth.NotFoundError('not found')
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock_urlfetch([
        ({
            'deadline': 10,
            'headers': {
                # Switches to own token.
                'Authorization': 'Bearer own-token',
            },
            'method': 'GET',
            'url': 'https://external.example.com/123',
        }, Response(200, '', {})),
    ])
    net.request(url='https://external.example.com/123', project_id='project-id') 

def test_cache_with_tasklets(self):
    @utils.cache
    def f():
      ndb.sleep(0).wait()  # Yield thread.
      return 1

    @ndb.tasklet
    def g():
      yield ()  # Make g a generator.
      raise ndb.Return(f())

    def test():
      ndb.Future.wait_all([(g()), (g())])

    t = threading.Thread(target=test)
    t.daemon = True
    t.start()
    t.join(1)
    if t.is_alive():
      self.fail('deadlock') 

def mock_methods(self):
    calls = []

    @ndb.tasklet
    def via_jwt(*args):
      calls.append(('jwt_based', args))
      raise ndb.Return({'access_token': 'token', 'exp_ts': 0})

    self.mock(service_account, '_mint_jwt_based_token_async', via_jwt)

    def via_gae_api(*args):
      calls.append(('gae_api', args))
      return 'token', 0
    self.mock(service_account.app_identity, 'get_access_token', via_gae_api)

    return calls 

def setUp(self):
    super(NetTest, self).setUp()

    @ndb.tasklet
    def get_access_token(*_args):
      raise ndb.Return(('own-token', 0))
    self.mock(auth, 'get_access_token_async', get_access_token)

    @ndb.tasklet
    def get_project_access_token(project_id, _scopes):
      raise ndb.Return(('%s-token' % project_id, 0))
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock(
        auth, 'is_internal_domain',
        lambda domain: domain == 'internal.example.com')

    self.mock(logging, 'warning', lambda *_args: None)
    self.mock(logging, 'error', lambda *_args: None) 

def test_project_tokens_fallback(self):
    @ndb.tasklet
    def get_project_access_token(_project_id, _scopes):
      raise auth.NotFoundError('not found')
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock_urlfetch([
        ({
            'deadline': 10,
            'headers': {
                # Switches to own token.
                'Authorization': 'Bearer own-token',
            },
            'method': 'GET',
            'url': 'https://external.example.com/123',
        }, Response(200, '', {})),
    ])
    net.request(url='https://external.example.com/123', project_id='project-id') 

def _get_project_configs_async(project_ids, path, message_factory):
  """Returns a mapping {project_id: message}.

  If a project does not exist, the message is None.
  """
  assert isinstance(project_ids, list)
  if not project_ids:
    empty = ndb.Future()
    empty.set_result({})
    return empty

  @ndb.tasklet
  def get_async():
    prefix = 'projects/'
    msgs = yield storage.get_latest_messages_async(
        [prefix + pid for pid in _filter_existing(project_ids)],
        path, message_factory)
    raise ndb.Return({
      # msgs may not have a key because we filter project ids by existence
      pid: msgs.get(prefix + pid)
      for pid in project_ids
    })

  return get_async() 

def test_unordered_first_concurrent_jobs(self):
    items = [10, 5, 0]

    log = []

    @ndb.tasklet
    def fn_async(x):
      log.append('%d started' % x)
      yield ndb.sleep(float(x) / 1000)
      log.append('%d finishing' % x)
      raise ndb.Return(x)

    expected = [(5, 5), (0, 0), (10, 10)]
    actual = utils.async_apply(
        items, fn_async, concurrent_jobs=2, unordered=True)
    self.assertFalse(isinstance(actual, list))
    self.assertEqual(expected, list(actual))
    self.assertEqual(log, [
        '10 started',
        '5 started',
        '5 finishing',
        '0 started',
        '0 finishing',
        '10 finishing',
    ]) 

def mock_methods(self):
    calls = []

    @ndb.tasklet
    def via_jwt(*args):
      calls.append(('jwt_based', args))
      raise ndb.Return({'access_token': 'token', 'exp_ts': 0})

    self.mock(service_account, '_mint_jwt_based_token_async', via_jwt)

    def via_gae_api(*args):
      calls.append(('gae_api', args))
      return 'token', 0
    self.mock(service_account.app_identity, 'get_access_token', via_gae_api)

    return calls 

def test_project_tokens_fallback(self):
    @ndb.tasklet
    def get_project_access_token(_project_id, _scopes):
      raise auth.NotFoundError('not found')
    self.mock(auth, 'get_project_access_token_async', get_project_access_token)

    self.mock_urlfetch([
        ({
            'deadline': 10,
            'headers': {
                # Switches to own token.
                'Authorization': 'Bearer own-token',
            },
            'method': 'GET',
            'url': 'https://external.example.com/123',
        }, Response(200, '', {})),
    ])
    net.request(url='https://external.example.com/123', project_id='project-id') 

def test_cache_with_tasklets(self):
    @utils.cache
    def f():
      ndb.sleep(0).wait()  # Yield thread.
      return 1

    @ndb.tasklet
    def g():
      yield ()  # Make g a generator.
      raise ndb.Return(f())

    def test():
      ndb.Future.wait_all([(g()), (g())])

    t = threading.Thread(target=test)
    t.daemon = True
    t.start()
    t.join(1)
    if t.is_alive():
      self.fail('deadlock') 

def _check_etag(self, etag):
    """Check if etag is the same across requests to GCS.

    If self._etag is None, set it. If etag is set, check that the new
    etag equals the old one.

    In the __init__ method, we fire one HEAD and one GET request using
    ndb tasklet. One of them would return first and set the first value.

    Args:
      etag: etag from a GCS HTTP response. None if etag is not part of the
        response header. It could be None for example in the case of GCS
        composite file.

    Raises:
      ValueError: if two etags are not equal.
    """
    if etag is None:
      return
    elif self._etag is None:
      self._etag = etag
    elif self._etag != etag:
      raise ValueError('File on GCS has changed while reading.') 

def define_get_google():
    @ndb.tasklet
    def get_google():
        context = ndb.get_context()
        result = yield context.urlfetch("http://www.google.com/")
        if result.status_code == 200:
            raise ndb.Return(result.content)
        # else return None

    return get_google 

def test_mint_jwt_based_token(self):
    self.mock_now(datetime.datetime(2015, 1, 2, 3))
    self.mock(os, 'urandom', lambda x: '1'*x)

    calls = []
    @ndb.tasklet
    def mocked_call(**kwargs):
      calls.append(kwargs)
      raise ndb.Return({'access_token': 'token', 'expires_in': 3600})
    self.mock(service_account, '_call_async', mocked_call)

    signer = FakeSigner()
    token = service_account._mint_jwt_based_token_async(
        ['scope1', 'scope2'], signer).get_result()
    self.assertEqual({'access_token': 'token', 'exp_ts': 1420171200.0}, token)

    self.assertEqual([{
      'aud': 'https://www.googleapis.com/oauth2/v4/token',
      'exp': 1420171195,
      'iat': 1420167595,
      'iss': 'fake@example.com',
      'jti': 'MTExMTExMTExMTExMTExMQ',
      'scope': 'scope1 scope2',
    }], signer.claimsets)

    self.assertEqual([
      {
        'url': 'https://www.googleapis.com/oauth2/v4/token',
        'method': 'POST',
        'headers': {
          'Accept': 'application/json',
          'Content-Type': 'application/x-www-form-urlencoded',
        },
        'payload':
            'grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&'
            'assertion=fake_jwt',
      }], calls) 

def setUp(self):
    super(GerritFetchTestCase, self).setUp()
    self.response = mock.Mock(
        status_code=httplib.OK,
        headers={},
        content='',
    )
    self.urlfetch_mock = mock.Mock(return_value=self.response)
    @ndb.tasklet
    def mocked_urlfetch(**kwargs):
      raise ndb.Return(self.urlfetch_mock(**kwargs))
    self.mock(net, 'urlfetch_async', mocked_urlfetch)
    self.mock(auth, 'get_access_token', mock.Mock(return_value=('token', 0.0)))
    self.mock(logging, 'warning', mock.Mock())