Browse Source

[1828] Make code safer by using with

Mukund Sivaraman 13 years ago
parent
commit
1ece2eadd4
1 changed files with 17 additions and 14 deletions
  1. 17 14
      src/bin/bindctl/tests/bindctl_test.py

+ 17 - 14
src/bin/bindctl/tests/bindctl_test.py

@@ -425,6 +425,12 @@ class FakeBindCmdInterpreter(bindcmd.BindCmdInterpreter):
 
 class TestBindCmdInterpreter(unittest.TestCase):
 
+    def setUp(self):
+        self.old_stdout = sys.stdout
+
+    def tearDown(self):
+        sys.stdout = self.old_stdout
+
     def _create_invalid_csv_file(self, csvfilename):
         import csv
         csvfile = open(csvfilename, 'w')
@@ -447,20 +453,17 @@ class TestBindCmdInterpreter(unittest.TestCase):
         self.assertEqual(new_csv_dir, custom_cmd.csv_file_dir)
 
     def test_get_saved_user_info(self):
-        old_stdout = sys.stdout
-        sys.stdout = open(os.devnull, 'w')
-        cmd = bindcmd.BindCmdInterpreter()
-        users = cmd._get_saved_user_info('/notexist', 'csv_file.csv')
-        self.assertEqual([], users)
-
-        csvfilename = 'csv_file.csv'
-        self._create_invalid_csv_file(csvfilename)
-        users = cmd._get_saved_user_info('./', csvfilename)
-        self.assertEqual([], users)
-        os.remove(csvfilename)
-        sys.stdout.close()
-        sys.stdout = old_stdout
-
+        with open(os.devnull, 'w') as f:
+            sys.stdout = f
+            cmd = bindcmd.BindCmdInterpreter()
+            users = cmd._get_saved_user_info('/notexist', 'csv_file.csv')
+            self.assertEqual([], users)
+
+            csvfilename = 'csv_file.csv'
+            self._create_invalid_csv_file(csvfilename)
+            users = cmd._get_saved_user_info('./', csvfilename)
+            self.assertEqual([], users)
+            os.remove(csvfilename)
 
 class TestCommandLineOptions(unittest.TestCase):
     def setUp(self):