How can I speed up or perftest context management in Google cloud-ndb library?

Question

I'm in the process of migrating an app running in the Google App Engine Python 2.7 Standard Environment from the built-in NDB to library to the newer cloud-ndb library (prior to a migration to Python 3). The new library requires context for all NDB calls, and after some research and conversation with others on an appengine email list I ran with an approach that enters the context in unittest setUp() and tearDown(). All of my test files derive from a shim base test class.

#speed_test.py
from test._base_test import MyTestCase

class SpeedTest(MyTestCase):
  def test_simple(self):
    string = 'string'
    self.assertEqual('string', string)

#_base_test.py
import requests
import unittest
import webapp2

from google.cloud import ndb


class MyTestCase(unittest.TestCase):
  def setUp(self):
    self.testapp = webapp2.WSGIApplication([])
    # clear datastore
    requests.post('http://localhost:8089/reset')
    self.ndb_context = ndb.Client().context()
    self.ndb_context.__enter__()

  def tearDown(self):
    self.ndb_context.__exit__(None, None, None)
    mock.patch.stopall()

The problem is that entering the context this way adds a little under 2 seconds to each one of my tests. That's mildly irritating for my test files with a dozen or fewer tests, but our full test suite has 681 tests and this speed just isn't feasible.

The above simple speed test should take a fraction of a second to complete, but it consistently runs in 1.7s or more. If I comment out self.ndb_context = ndb.Client().context() and self.ndb_context.__enter__(), it drops down to about 0.003s.

This feels like a long time for the context manager to run but I'm not sure how to do debug performance further at this point. I'm happy to take suggestions on a better approach or means to debug what's happening here so I can file an issue w/ the library if needed. Thanks!

Answer 1

The issue turned out to be with creating the NDB client object repeatedly in the setUp()/tearDown() methods. Once I moved the client to a variable outside the class it returned to the expected levels of performance.

#_base_test.py
import requests
import unittest
import webapp2

from google.cloud import ndb

NDB_CLIENT = ndb.Client()

class MyTestCase(unittest.TestCase):
  def setUp(self):
    self.testapp = webapp2.WSGIApplication([])
    # clear datastore
    requests.post('http://localhost:8089/reset')
    self.ndb_context = NDB_CLIENT.context()
    self.ndb_context.__enter__()

  def tearDown(self):
    self.ndb_context.__exit__(None, None, None)
    mock.patch.stopall()