# HG changeset patch
# User Adrien Di Mascio <Adrien.DiMascio@logilab.fr>
# Date 1282909495 -7200
# Node ID 7da7c270bbbdd454612b6b1951fbe87b182e58b4
# Parent  ceb6951f9d2f5d286e86ef39e5a03c9a9f5f8542
[devtools] make CubiwebTC.login() context-manager restore the testcase connection on exit

diff -r ceb6951f9d2f -r 7da7c270bbbd devtools/testlib.py
--- a/devtools/testlib.py	Fri Aug 27 09:02:41 2010 +0200
+++ b/devtools/testlib.py	Fri Aug 27 13:44:55 2010 +0200
@@ -142,6 +142,29 @@
 
 cwconfig.SMTP = MockSMTP
 
+class TestCaseConnectionProxy(object):
+    """thin wrapper around `cubicweb.dbapi.Connection` context-manager
+    used in CubicWebTC (cf. `cubicweb.devtools.testlib.CubicWebTC.login` method)
+
+    It just proxies to the default connection context manager but
+    restores the original connection on exit.
+    """
+    def __init__(self, testcase, cnx):
+        self.testcase = testcase
+        self.cnx = cnx
+
+    def __getattr__(self, attrname):
+        return getattr(self.cnx, attrname)
+
+    def __enter__(self):
+        return self.cnx.__enter__()
+
+    def __exit__(self, exctype, exc, tb):
+        try:
+            return self.cnx.__exit__(exctype, exc, tb)
+        finally:
+            self.cnx.close()
+            self.testcase.restore_connection()
 
 # base class for cubicweb tests requiring a full cw environments ###############
 
@@ -331,7 +354,7 @@
             self._cnxs.append(self.cnx)
         if login == self.vreg.config.anonymous_user()[0]:
             self.cnx.anonymous_connection = True
-        return self.cnx
+        return TestCaseConnectionProxy(self, self.cnx)
 
     def restore_connection(self):
         if not self.cnx is self._orig_cnx[0]: