Use a decorator to save and restore thread-local state.

lib/python/Testing/ZopeTestCase/functional.py

@@ -23,6 +23,25 @@ import sandbox
 import interfaces
 
 
+def savestate(func):
+    '''Decorator saving thread local state before executing func
+       and restoring it afterwards.
+    '''
+    from AccessControl.SecurityManagement import getSecurityManager
+    from AccessControl.SecurityManagement import setSecurityManager
+    from zope.app.component.hooks import getSite
+    from zope.app.component.hooks import setSite
+
+    def wrapped_func(*args, **kw):
+        sm, site = getSecurityManager(), getSite()
+        try:
+            return func(*args, **kw)
+        finally:
+            setSecurityManager(sm)
+            setSite(site)
+    return wrapped_func
+
+
 class Functional(sandbox.Sandboxed):
     '''Derive from this class and an xTestCase to get functional
        testing support::
@@ -33,25 +52,15 @@ class Functional(sandbox.Sandboxed):
 
     __implements__ = (interfaces.IFunctional,)
 
+    @savestate
     def publish(self, path, basic=None, env=None, extra=None,
                 request_method='GET', stdin=None, handle_errors=True):
         '''Publishes the object at 'path' returning a response object.'''
 
-        from zope.app.component.hooks import setSite, getSite
         from StringIO import StringIO
         from ZPublisher.Response import Response
         from ZPublisher.Test import publish_module
 
-        from AccessControl.SecurityManagement import getSecurityManager
-        from AccessControl.SecurityManagement import setSecurityManager
-
-        # Save current security manager
-        sm = getSecurityManager()
-
-        # And we need to store the old site
-        old_site = getSite()
-        setSite(None)
-
         # Commit the sandbox for good measure
         transaction.commit()
 
@@ -91,12 +100,6 @@ class Functional(sandbox.Sandboxed):
                        debug=not handle_errors,
                       )
 
-        # Restore security manager
-        setSecurityManager(sm)
-
-        # And we need to restore the site again
-        setSite(old_site)
-
         return ResponseWrapper(response, outstream, path)
 
 

lib/python/Testing/ZopeTestCase/zopedoctest/functional.py

@@ -31,6 +31,7 @@ from Testing.ZopeTestCase import user_role
 from Testing.ZopeTestCase import standard_permissions
 from Testing.ZopeTestCase.sandbox import AppZapper
 from Testing.ZopeTestCase.functional import ResponseWrapper
+from Testing.ZopeTestCase.functional import savestate
 
 
 class HTTPHeaderOutput:
@@ -110,6 +111,7 @@ def sync():
     getRootFolder()._p_jar.sync()
 
 
+@savestate
 def http(request_string, handle_errors=True):
     """Execute an HTTP request string via the publisher
 
@@ -117,19 +119,9 @@ def http(request_string, handle_errors=True):
     """
     import urllib
     import rfc822
-    from zope.app.component.hooks import setSite, getSite
     from cStringIO import StringIO
     from ZPublisher.Response import Response
     from ZPublisher.Test import publish_module
-    from AccessControl.SecurityManagement import getSecurityManager
-    from AccessControl.SecurityManagement import setSecurityManager
-
-    # Save current Security Manager
-    old_sm = getSecurityManager()
-
-    # And we need to store the old site
-    old_site = getSite()
-    setSite(None)
 
     # Commit work done by previous python code.
     transaction.commit()
@@ -194,14 +186,6 @@ def http(request_string, handle_errors=True):
     header_output.appendResponseHeaders(response._cookie_list())
     header_output.appendResponseHeaders(response.accumulated_headers.splitlines())
 
-    # Restore previous security manager, which may have been changed
-    # by calling the publish method above
-    setSecurityManager(old_sm)
-
-    # And we need to restore the site again
-    setSite(old_site)
-    # Sync connection
-
     sync()
 
     return DocResponseWrapper(response, outstream, path, header_output)