Handle failure in setup of ConnectionTwoPhaseTest

Previously, this test had a bug, because if the CREATE TABLE statement
failed, the setup would fail without committing or rolling back the
active transaction.

1 files changed, 5 insertions, 3 deletions

diff --git a/tests/test_connection.py b/tests/test_connection.py
index f1668d5..8236801 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -913,9 +913,11 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
             cur.execute("DROP TABLE test_tpc;")
         except psycopg2.ProgrammingError:
             cnn.rollback()
-        cur.execute("CREATE TABLE test_tpc (data text);")
-        cnn.commit()
-        cnn.close()
+        try:
+            cur.execute("CREATE TABLE test_tpc (data text);")
+            cnn.commit()
+        finally:
+            cnn.close()
 
     def count_xacts(self):
         """Return the number of prepared xacts currently in the test db."""