minor improvements

Author: funbringer

testgres/utils.py

@@ -38,7 +38,7 @@ def release_port(port):
     Free port provided by reserve_port().
     """
 
-    bound_ports.remove(port)
+    bound_ports.discard(port)
 
 
 def execute_utility(args, logfile=None):

tests/test_simple.py

@@ -11,6 +11,7 @@
 
 import logging.config
 
+from contextlib import contextmanager
 from distutils.version import LooseVersion
 
 from testgres import \
@@ -58,6 +59,15 @@ def good_properties(f):
             return True
 
 
+@contextmanager
+def removing(f):
+    try:
+        yield f
+    finally:
+        if os.path.isfile(f):
+            os.remove(f)
+
+
 class SimpleTest(unittest.TestCase):
     def test_custom_init(self):
         with get_new_node() as node:
@@ -382,19 +392,14 @@ def test_dump(self):
             node1.execute(query_create)
 
             # take a new dump
-            dump = node1.dump()
-            self.assertTrue(os.path.isfile(dump))
-
-            with get_new_node().init().start() as node2:
-
-                # restore dump
-                node2.restore(filename=dump)
-
-                res = node2.execute(query_select)
-                self.assertListEqual(res, [(1, ), (2, )])
-
-            # finally, remove dump
-            os.remove(dump)
+            with removing(node1.dump()) as dump:
+                with get_new_node().init().start() as node2:
+                    # restore dump
+                    self.assertTrue(os.path.isfile(dump))
+                    node2.restore(filename=dump)
+
+                    res = node2.execute(query_select)
+                    self.assertListEqual(res, [(1, ), (2, )])
 
     def test_users(self):
         with get_new_node().init().start() as node: