Replace self.mock_* pairs by self.mocked_* context managers

Make sure cleanup happens

Author: Flamefire

easybuild/base/testing.py

@@ -206,30 +206,30 @@ def get_stderr(self):
         return sys.stderr.getvalue()
 
     @contextmanager
-    def mocked_stdout(self):
+    def mocked_stdout(self, force_tty=False):
         """Context manager to mock stdout"""
         self.mock_stdout(True)
         try:
             yield sys.stdout
         finally:
-            self.mock_stdout(False)
+            self.mock_stdout(False, force_tty=force_tty)
 
     @contextmanager
-    def mocked_stderr(self):
+    def mocked_stderr(self, force_tty=False):
         """Context manager to mock stdout"""
         self.mock_stderr(True)
         try:
             yield sys.stderr
         finally:
-            self.mock_stderr(False)
+            self.mock_stderr(False, force_tty=force_tty)
 
     @contextmanager
-    def mocked_stdout_stderr(self, mock_stdout=True, mock_stderr=True):
+    def mocked_stdout_stderr(self, mock_stdout=True, mock_stderr=True, force_tty=False):
         """Context manager to mock stdout and stderr"""
         if mock_stdout:
-            self.mock_stdout(True)
+            self.mock_stdout(True, force_tty=force_tty)
         if mock_stderr:
-            self.mock_stderr(True)
+            self.mock_stderr(True, force_tty=force_tty)
         try:
             if mock_stdout and mock_stderr:
                 yield sys.stdout, sys.stderr

test/framework/build_log.py

@@ -111,19 +111,18 @@ def test_easybuildlog(self):
         # test basic log methods
         logToFile(tmplog, enable=True)
         log = getLogger('test_easybuildlog')
-        self.mock_stderr(True)
-        log.setLevelName('DEBUG')
-        log.debug("123 debug")
-        log.info("foobar info")
-        log.warning("justawarning")
-        log.deprecated("anotherwarning", newer_ver)
-        log.deprecated("onemorewarning", '1.0', '2.0')
-        log.deprecated("lastwarning", '1.0', max_ver='2.0')
-        log.deprecated("thisisnotprinted", '1.0', max_ver='2.0', silent=True)
-        log.error("kaput")
-        log.error("err: %s", 'msg: %s')
-        stderr = self.get_stderr()
-        self.mock_stderr(False)
+        with self.mocked_stderr():
+            log.setLevelName('DEBUG')
+            log.debug("123 debug")
+            log.info("foobar info")
+            log.warning("justawarning")
+            log.deprecated("anotherwarning", newer_ver)
+            log.deprecated("onemorewarning", '1.0', '2.0')
+            log.deprecated("lastwarning", '1.0', max_ver='2.0')
+            log.deprecated("thisisnotprinted", '1.0', max_ver='2.0', silent=True)
+            log.error("kaput")
+            log.error("err: %s", 'msg: %s')
+            stderr = self.get_stderr()
 
         more_info = "see https://docs.easybuild.io/deprecated-functionality/ for more information"
         common_warning = "\nWARNING: Deprecated functionality, will no longer work in"
@@ -188,9 +187,8 @@ def test_easybuildlog(self):
         logToFile(tmplog, enable=True)
 
         # also test use of 'more_info' named argument for log.deprecated
-        self.mock_stderr(True)
-        log.deprecated("\nthis is just a test\n", newer_ver, more_info="(see URLGOESHERE for more information)")
-        self.mock_stderr(False)
+        with self.mocked_stderr():
+            log.deprecated("\nthis is just a test\n", newer_ver, more_info="(see URLGOESHERE for more information)")
         logtxt = read_file(tmplog)
         expected_logtxt = '\n'.join([
             "[WARNING] :: Deprecated functionality, will no longer work in EasyBuild v10000001: ",
@@ -211,18 +209,17 @@ def test_log_levels(self):
         logToFile(tmplog, enable=True)
         log = getLogger('test_easybuildlog')
 
-        self.mock_stderr(True)  # avoid that some log statement spit out stuff to stderr while tests are running
-        for level in ['ERROR', 'WARNING', 'INFO', 'DEBUG', 'DEVEL']:
-            log.setLevelName(level)
-            log.raiseError = False
-            log.error('kaput')
-            log.deprecated('almost kaput', '10000000000000')
-            log.raiseError = True
-            log.warning('this is a warning')
-            log.info('fyi')
-            log.debug('gdb')
-            log.devel('tmi')
-        self.mock_stderr(False)
+        with self.mocked_stderr():  # avoid that some log statement spit out stuff to stderr while tests are running
+            for level in ['ERROR', 'WARNING', 'INFO', 'DEBUG', 'DEVEL']:
+                log.setLevelName(level)
+                log.raiseError = False
+                log.error('kaput')
+                log.deprecated('almost kaput', '10000000000000')
+                log.raiseError = True
+                log.warning('this is a warning')
+                log.info('fyi')
+                log.debug('gdb')
+                log.devel('tmi')
 
         logToFile(tmplog, enable=False)
         logtxt = read_file(tmplog)
@@ -249,13 +246,10 @@ def test_print_warning(self):
         """Test print_warning"""
         def run_check(args, silent=False, expected_stderr='', **kwargs):
             """Helper function to check stdout/stderr produced via print_warning."""
-            self.mock_stderr(True)
-            self.mock_stdout(True)
-            print_warning(*args, silent=silent, **kwargs)
-            stderr = self.get_stderr()
-            stdout = self.get_stdout()
-            self.mock_stdout(False)
-            self.mock_stderr(False)
+            with self.mocked_stdout_stderr():
+                print_warning(*args, silent=silent, **kwargs)
+                stderr = self.get_stderr()
+                stdout = self.get_stdout()
             self.assertEqual(stdout, '')
             self.assertEqual(stderr, expected_stderr)
 
@@ -308,13 +302,10 @@ def test_print_msg(self):
         """Test print_msg"""
         def run_check(msg, args, expected_stdout='', expected_stderr='', **kwargs):
             """Helper function to check stdout/stderr produced via print_msg."""
-            self.mock_stdout(True)
-            self.mock_stderr(True)
-            print_msg(msg, *args, **kwargs)
-            stdout = self.get_stdout()
-            stderr = self.get_stderr()
-            self.mock_stdout(False)
-            self.mock_stderr(False)
+            with self.mocked_stdout_stderr():
+                print_msg(msg, *args, **kwargs)
+                stdout = self.get_stdout()
+                stderr = self.get_stderr()
             self.assertEqual(stdout, expected_stdout)
             self.assertEqual(stderr, expected_stderr)
 
@@ -347,13 +338,10 @@ def test_print_msg_rich(self):
 
         def run_check(msg, args, expected_stdout='', expected_stderr='', **kwargs):
             """Helper function to check stdout/stderr produced via print_msg."""
-            self.mock_stdout(True, force_tty=True)
-            self.mock_stderr(True, force_tty=True)
-            print_msg(msg, *args, **kwargs)
-            stdout = self.get_stdout()
-            stderr = self.get_stderr()
-            self.mock_stdout(False)
-            self.mock_stderr(False)
+            with self.mocked_stdout_stderr(force_tty=True):
+                print_msg(msg, *args, **kwargs)
+                stdout = self.get_stdout()
+                stderr = self.get_stderr()
             # Check that the expected output stdout/stderr is being colorized
             # In order to allow using more styles/schemes (or even to account to changes in the rich default)
             # We only check for \x1b[0m which is the ANSI reset code that should be present in the output to undo
@@ -386,13 +374,10 @@ def test_dry_run_msg(self):
         """Test dry_run_msg"""
         def run_check(msg, args, expected_stdout='', **kwargs):
             """Helper function to check stdout/stderr produced via dry_run_msg."""
-            self.mock_stdout(True)
-            self.mock_stderr(True)
-            dry_run_msg(msg, *args, **kwargs)
-            stdout = self.get_stdout()
-            stderr = self.get_stderr()
-            self.mock_stdout(False)
-            self.mock_stderr(False)
+            with self.mocked_stdout_stderr():
+                dry_run_msg(msg, *args, **kwargs)
+                stdout = self.get_stdout()
+                stderr = self.get_stderr()
             self.assertEqual(stdout, expected_stdout)
             self.assertEqual(stderr, '')
 
@@ -407,13 +392,10 @@ def test_dry_run_warning(self):
         """Test dry_run_warningmsg"""
         def run_check(msg, args, expected_stdout='', **kwargs):
             """Helper function to check stdout/stderr produced via dry_run_warningmsg."""
-            self.mock_stdout(True)
-            self.mock_stderr(True)
-            dry_run_warning(msg, *args, **kwargs)
-            stdout = self.get_stdout()
-            stderr = self.get_stderr()
-            self.mock_stdout(False)
-            self.mock_stderr(False)
+            with self.mocked_stdout_stderr():
+                dry_run_warning(msg, *args, **kwargs)
+                stdout = self.get_stdout()
+                stderr = self.get_stderr()
             self.assertEqual(stdout, expected_stdout)
             self.assertEqual(stderr, '')
 
@@ -458,10 +440,9 @@ def test_init_logging(self):
         stop_logging(logfile)
 
         # by default, path to tmp log file is printed
-        self.mock_stdout(True)
-        log, logfile = init_logging(None)
-        stdout = self.get_stdout()
-        self.mock_stdout(False)
+        with self.mocked_stdout():
+            log, logfile = init_logging(None)
+            stdout = self.get_stdout()
         self.assertExists(logfile)
         self.assertEqual(os.path.dirname(logfile), tmpdir)
         self.assertIsInstance(log, EasyBuildLog)
@@ -470,9 +451,8 @@ def test_init_logging(self):
         stop_logging(logfile)
 
         # logging to stdout implies no log file
-        self.mock_stdout(True)
-        log, logfile = init_logging(None, logtostdout=True)
-        self.mock_stdout(False)
+        with self.mocked_stdout():
+            log, logfile = init_logging(None, logtostdout=True)
         self.assertEqual(logfile, None)
         self.assertIsInstance(log, EasyBuildLog)
 

test/framework/containers.py

@@ -74,13 +74,10 @@ class ContainersTest(EnhancedTestCase):
 
     def run_main(self, args, raise_error=True):
         """Helper function to run main with arguments specified in 'args' and return stdout/stderr."""
-        self.mock_stdout(True)
-        self.mock_stderr(True)
-        self.eb_main(args, raise_error=raise_error, verbose=True, do_build=True)
-        stdout = self.get_stdout().strip()
-        stderr = self.get_stderr().strip()
-        self.mock_stdout(False)
-        self.mock_stderr(False)
+        with self.mocked_stdout_stderr():
+            self.eb_main(args, raise_error=raise_error, verbose=True, do_build=True)
+            stdout = self.get_stdout().strip()
+            stderr = self.get_stderr().strip()
 
         return stdout, stderr
 
@@ -337,9 +334,8 @@ def test_end2end_singularity_image(self):
         # test again with container image already existing
 
         error_pattern = "Container image already exists at %s, not overwriting it without --force" % cont_img
-        self.mock_stdout(True)
-        self.assertErrorRegex(EasyBuildError, error_pattern, self.run_main, args, raise_error=True)
-        self.mock_stdout(False)
+        with self.mocked_stdout():
+            self.assertErrorRegex(EasyBuildError, error_pattern, self.run_main, args, raise_error=True)
 
         args.append('--force')
         stdout, stderr = self.run_main(args)
@@ -397,11 +393,12 @@ def test_end2end_dockerfile(self):
 
         error_pattern = "Container recipe at %s/containers/Dockerfile.toy-0.0 already exists, " \
                         "not overwriting it without --force" % self.test_prefix
-        self.assertErrorRegex(EasyBuildError,
-                              error_pattern,
-                              self.run_main,
-                              base_args + ['--container-config=centos:7'],
-                              raise_error=True)
+        with self.mocked_stdout():
+            self.assertErrorRegex(EasyBuildError,
+                                  error_pattern,
+                                  self.run_main,
+                                  base_args + ['--container-config=centos:7'],
+                                  raise_error=True)
 
         remove_file(os.path.join(self.test_prefix, 'containers', 'Dockerfile.toy-0.0'))
 

test/framework/docs.py

@@ -1349,13 +1349,10 @@ def test_help(self):
         Test output produced by --help, with various output formats
         """
         def get_eb_help_output(arg=''):
-            self.mock_stderr(True)
-            self.mock_stdout(True)
-            self.eb_main(['--help', arg])
-            stderr = self.get_stderr()
-            stdout = self.get_stdout()
-            self.mock_stderr(False)
-            self.mock_stdout(False)
+            with self.mocked_stdout_stderr():
+                self.eb_main(['--help', arg])
+                stderr = self.get_stderr()
+                stdout = self.get_stdout()
 
             self.assertFalse(stderr)
             return stdout