Commits

Anonymous committed 57ea671

switched to new testbed feature for improved unit test performance

Comments (0)

Files changed (4)

-from ..utils import appid, have_appserver, on_production_server
 from ..boot import DATA_ROOT
+from ..utils import appid, on_production_server
 from .creation import DatabaseCreation
+from .stubs import stub_manager
 from django.db.backends.util import format_number
 from djangotoolbox.db.base import NonrelDatabaseFeatures, \
     NonrelDatabaseOperations, NonrelDatabaseWrapper, NonrelDatabaseClient, \
 from google.appengine.ext.db.metadata import get_kinds, get_namespaces
 from google.appengine.api.datastore import Query, Delete
 from google.appengine.api.namespace_manager import set_namespace
-from urllib2 import HTTPError, URLError
 import logging
 import os
-import time
 
-REMOTE_API_SCRIPT = '$PYTHON_LIB/google/appengine/ext/remote_api/handler.py'
 DATASTORE_PATHS = {
     'datastore_path': os.path.join(DATA_ROOT, 'datastore'),
     'blobstore_path': os.path.join(DATA_ROOT, 'blobstore'),
     'prospective_search_path': os.path.join(DATA_ROOT, 'prospective-search'),
 }
 
-def auth_func():
-    import getpass
-    return raw_input('Login via Google Account (see note above if login fails): '), getpass.getpass('Password: ')
-
-def rpc_server_factory(*args, ** kwargs):
-    from google.appengine.tools import appengine_rpc
-    kwargs['save_cookies'] = True
-    return appengine_rpc.HttpRpcServer(*args, ** kwargs)
-
 def get_datastore_paths(options):
     paths = {}
     for key, path in DATASTORE_PATHS.items():
         paths[key] = options.get(key, path)
     return paths
 
-def get_test_datastore_paths(options, inmemory=True):
-    paths = get_datastore_paths(options)
-    for key in paths:
-        paths[key] += '.test'
-    if inmemory:
-        for key in ('datastore_path', 'blobstore_path'):
-            paths[key] = None
-    return paths
-
 def destroy_datastore(paths):
     """Destroys the appengine datastore at the specified paths."""
     for path in paths.values():
         self.validation = DatabaseValidation(self)
         self.introspection = DatabaseIntrospection(self)
         options = self.settings_dict
-        self.use_test_datastore = False
-        self.test_datastore_inmemory = True
-        self.remote = options.get('REMOTE', False)
-        if on_production_server:
-            self.remote = False
         self.remote_app_id = options.get('REMOTE_APP_ID', appid)
         self.domain = options.get('DOMAIN', 'appspot.com')
         self.remote_api_path = options.get('REMOTE_API_PATH', None)
         self.secure_remote_api = options.get('SECURE_REMOTE_API', True)
-        self._setup_stubs()
 
-    def _get_paths(self):
-        if self.use_test_datastore:
-            return get_test_datastore_paths(self.settings_dict, self.test_datastore_inmemory)
+        remote = options.get('REMOTE', False)
+        if on_production_server:
+            remote = False
+        if remote:
+            stub_manager.setup_remote_stubs(self)
         else:
-            return get_datastore_paths(self.settings_dict)
-
-    def _setup_stubs(self):
-        # If this code is being run without an appserver (eg. via a django
-        # commandline flag) then setup a default stub environment.
-        if not have_appserver:
-            from google.appengine.tools import dev_appserver_main
-            args = dev_appserver_main.DEFAULT_ARGS.copy()
-            args.update(self._get_paths())
-            log_level = logging.getLogger().getEffectiveLevel()
-            logging.getLogger().setLevel(logging.WARNING)
-            from google.appengine.tools import dev_appserver
-            dev_appserver.SetupStubs(appid, **args)
-            logging.getLogger().setLevel(log_level)
-        # If we're supposed to set up the remote_api, do that now.
-        if self.remote:
-            self.setup_remote()
-
-    def setup_remote(self):
-        if not self.remote_api_path:
-            from ..utils import appconfig
-            for handler in appconfig.handlers:
-                if handler.script == REMOTE_API_SCRIPT:
-                    self.remote_api_path = handler.url.split('(', 1)[0]
-                    break
-        self.remote = True
-        server = '%s.%s' % (self.remote_app_id, self.domain)
-        remote_url = 'https://%s%s' % (server, self.remote_api_path)
-        logging.info('Setting up remote_api for "%s" at %s' %
-                     (self.remote_app_id, remote_url))
-        if not have_appserver:
-            print('Connecting to remote_api handler.\n\n'
-                  'IMPORTANT: Check your login method settings in the '
-                  'App Engine Dashboard if you have problems logging in. '
-                  'Login is only supported for Google Accounts.\n')
-        from google.appengine.ext.remote_api import remote_api_stub
-        remote_api_stub.ConfigureRemoteApi(None,
-            self.remote_api_path, auth_func, servername=server,
-            secure=self.secure_remote_api,
-            rpc_server_factory=rpc_server_factory)
-        retry_delay = 1
-        while retry_delay <= 16:
-            try:
-                remote_api_stub.MaybeInvokeAuthentication()
-            except HTTPError, e:
-                if not have_appserver:
-                    print 'Retrying in %d seconds...' % retry_delay
-                time.sleep(retry_delay)
-                retry_delay *= 2
-            else:
-                break
-        else:
-            try:
-                remote_api_stub.MaybeInvokeAuthentication()
-            except HTTPError, e:
-                raise URLError("%s\n"
-                               "Couldn't reach remote_api handler at %s.\n"
-                               "Make sure you've deployed your project and "
-                               "installed a remote_api handler in app.yaml. "
-                               "Note that login is only supported for "
-                               "Google Accounts. Make sure you've configured "
-                               "the correct authentication method in the "
-                               "App Engine Dashboard."
-                               % (e, remote_url))
-        logging.info('Now using the remote datastore for "%s" at %s' %
-                     (self.remote_app_id, remote_url))
+            stub_manager.setup_stubs(self)
 
     def flush(self):
         """Helper function to remove the current datastore and re-open the stubs"""
-        if self.remote:
+        if stub_manager.active_stubs == 'remote':
             import random
             import string
             code = ''.join([random.choice(string.ascii_letters) for x in range(4)])
             else:
                 print 'Aborting'
                 exit()
+        elif stub_manager.active_stubs == 'test':
+            stub_manager.deactivate_test_stubs()
+            stub_manager.activate_test_stubs()
 #        elif on_production_server or have_appserver:
 #            delete_all_entities()
         else:
-            destroy_datastore(self._get_paths())
-        self._setup_stubs()
+            destroy_datastore(get_datastore_paths(self.settings_dict))
+            stub_manager.setup_local_stubs(self)
 
 def delete_all_entities():
     for namespace in get_namespaces():
 from .db_settings import get_model_indexes
+from .stubs import stub_manager
 from djangotoolbox.db.creation import NonrelDatabaseCreation
 
 class StringType(object):
 
     data_types = get_data_types()
 
-    def create_test_db(self, *args, **kw):
-        """Destroys the test datastore. A new store will be recreated on demand"""
-        self.destroy_test_db()
-        self.connection.use_test_datastore = True
-        self.connection.flush()
+    def _create_test_db(self, *args, **kw):
+        self._had_test_stubs = stub_manager.active_stubs != 'test'
+        if self._had_test_stubs:
+            stub_manager.activate_test_stubs()
 
-    def destroy_test_db(self, *args, **kw):
-        """Destroys the test datastore files."""
-        from .base import destroy_datastore, get_test_datastore_paths
-        destroy_datastore(get_test_datastore_paths(self.connection.settings_dict))
+    def _destroy_test_db(self, *args, **kw):
+        if self._had_test_stubs:
+            stub_manager.deactivate_test_stubs()
+            stub_manager.setup_stubs(self.connection)
+        del self._had_test_stubs
+from ..utils import appid, have_appserver
+from google.appengine.ext.testbed import Testbed
+from urllib2 import HTTPError, URLError
+import logging
+import time
+
+REMOTE_API_SCRIPT = '$PYTHON_LIB/google/appengine/ext/remote_api/handler.py'
+
+def auth_func():
+    import getpass
+    return raw_input('Login via Google Account (see note above if login fails): '), getpass.getpass('Password: ')
+
+def rpc_server_factory(*args, ** kwargs):
+    from google.appengine.tools import appengine_rpc
+    kwargs['save_cookies'] = True
+    return appengine_rpc.HttpRpcServer(*args, ** kwargs)
+
+class StubManager(object):
+    def __init__(self):
+        self.testbed = Testbed()
+        self.active_stubs = None
+        self.pre_test_stubs = None
+
+    def setup_stubs(self, connection):
+        if self.active_stubs is not None:
+            return
+        if not have_appserver:
+            self.setup_local_stubs(connection)
+
+    def activate_test_stubs(self):
+        if self.active_stubs == 'test':
+            return
+        self.testbed.activate()
+        self.pre_test_stubs = self.active_stubs
+        self.active_stubs = 'test'
+        self.testbed.init_datastore_v3_stub()
+        self.testbed.init_memcache_stub()
+        self.testbed.init_taskqueue_stub()
+        self.testbed.init_urlfetch_stub()
+        self.testbed.init_user_stub()
+        self.testbed.init_xmpp_stub()
+        self.testbed.init_channel_stub()
+
+    def deactivate_test_stubs(self):
+        if self.active_stubs == 'test':
+            self.testbed.deactivate()
+            self.active_stubs = self.pre_test_stubs
+
+    def setup_local_stubs(self, connection):
+        if self.active_stubs == 'local':
+            return
+        from .base import get_datastore_paths
+        from google.appengine.tools import dev_appserver_main
+        args = dev_appserver_main.DEFAULT_ARGS.copy()
+        args.update(get_datastore_paths(connection.settings_dict))
+        log_level = logging.getLogger().getEffectiveLevel()
+        logging.getLogger().setLevel(logging.WARNING)
+        from google.appengine.tools import dev_appserver
+        dev_appserver.SetupStubs(appid, **args)
+        logging.getLogger().setLevel(log_level)
+        self.active_stubs = 'local'
+
+    def setup_remote_stubs(self, connection):
+        if self.active_stubs == 'remote':
+            return
+        if not connection.remote_api_path:
+            from ..utils import appconfig
+            for handler in appconfig.handlers:
+                if handler.script == REMOTE_API_SCRIPT:
+                    connection.remote_api_path = handler.url.split('(', 1)[0]
+                    break
+        server = '%s.%s' % (connection.remote_app_id, connection.domain)
+        remote_url = 'https://%s%s' % (server, connection.remote_api_path)
+        logging.info('Setting up remote_api for "%s" at %s' %
+                     (connection.remote_app_id, remote_url))
+        if not have_appserver:
+            print('Connecting to remote_api handler.\n\n'
+                  'IMPORTANT: Check your login method settings in the '
+                  'App Engine Dashboard if you have problems logging in. '
+                  'Login is only supported for Google Accounts.\n')
+        from google.appengine.ext.remote_api import remote_api_stub
+        remote_api_stub.ConfigureRemoteApi(None,
+            connection.remote_api_path, auth_func, servername=server,
+            secure=connection.secure_remote_api,
+            rpc_server_factory=rpc_server_factory)
+        retry_delay = 1
+        while retry_delay <= 16:
+            try:
+                remote_api_stub.MaybeInvokeAuthentication()
+            except HTTPError, e:
+                if not have_appserver:
+                    print 'Retrying in %d seconds...' % retry_delay
+                time.sleep(retry_delay)
+                retry_delay *= 2
+            else:
+                break
+        else:
+            try:
+                remote_api_stub.MaybeInvokeAuthentication()
+            except HTTPError, e:
+                raise URLError("%s\n"
+                               "Couldn't reach remote_api handler at %s.\n"
+                               "Make sure you've deployed your project and "
+                               "installed a remote_api handler in app.yaml. "
+                               "Note that login is only supported for "
+                               "Google Accounts. Make sure you've configured "
+                               "the correct authentication method in the "
+                               "App Engine Dashboard."
+                               % (e, remote_url))
+        logging.info('Now using the remote datastore for "%s" at %s' %
+                     (connection.remote_app_id, remote_url))
+        self.active_stubs = 'remote'
+
+stub_manager = StubManager()

management/commands/remote.py

 
     def run_from_argv(self, argv):
         from django.db import connections
+        from ...db.base import DatabaseWrapper
+        from ...db.stubs import stub_manager
         for connection in connections.all():
-            if hasattr(connection, 'setup_remote'):
-                connection.setup_remote()
+            if isinstance(connection, DatabaseWrapper):
+                stub_manager.setup_remote_stubs(connection)
+                break
         argv = argv[:1] + argv[2:]
         execute_from_command_line(argv)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.