1import unittest
2
3from google.appengine.ext import testbed
4
5from google.appengine.ext.ndb import model
6from google.appengine.ext.ndb import tasklets
7
8import webapp2
9
10
11def main():
12    unittest.main()
13
14
15def check_webob_version(minimum_version):
16    try:
17        # WebOb < 1.0 (App Engine SDK).
18        from webob.statusreasons import status_reasons
19        from webob.headerdict import HeaderDict
20        return 0.96 >= minimum_version
21    except ImportError:
22        # WebOb >= 1.0.
23        from webob.util import status_reasons
24        from webob.headers import ResponseHeaders
25        return 1.0 <= minimum_version
26
27
28class BaseTestCase(unittest.TestCase):
29
30    def setUp(self):
31        """Set up the test framework.
32
33        Service stubs are available for the following services:
34
35        - Datastore (use init_datastore_v3_stub)
36        - Memcache (use init_memcache_stub)
37        - Task Queue (use init_taskqueue_stub)
38        - Images (only for dev_appserver; use init_images_stub)
39        - URL fetch (use init_urlfetch_stub)
40        - User service (use init_user_stub)
41        - XMPP (use init_xmpp_stub)
42        """
43        # First, create an instance of the Testbed class.
44        self.testbed = testbed.Testbed()
45
46        # Then activate the testbed, which prepares the service stubs for use.
47        self.testbed.activate()
48
49        # To set custom env vars, pass them as kwargs *after* activate().
50        # self.setup_env()
51
52        # Next, declare which service stubs you want to use.
53        self.testbed.init_datastore_v3_stub()
54        self.testbed.init_memcache_stub()
55        self.testbed.init_user_stub()
56
57        # Only when testing ndb.
58        self.reset_kind_map()
59        self.setup_context_cache()
60
61    def tearDown(self):
62        # This restores the original stubs so that tests do not interfere
63        # with each other.
64        self.testbed.deactivate()
65        # Clear thread-local variables.
66        self.clear_globals()
67
68    def reset_kind_map(self):
69        model.Model._reset_kind_map()
70
71    def setup_context_cache(self):
72        """Set up the context cache.
73
74        We only need cache active when testing the cache, so the default
75        behavior is to disable it to avoid misleading test results. Override
76        this when needed.
77        """
78        ctx = tasklets.get_context()
79        ctx.set_cache_policy(False)
80        ctx.set_memcache_policy(False)
81
82    def clear_globals(self):
83        webapp2._local.__release_local__()
84
85    def register_model(self, name, cls):
86        model.Model._kind_map[name] = cls
87