diff --git a/gramps/cli/test/user_test.py b/gramps/cli/test/user_test.py index a02a751bc..050c5335c 100644 --- a/gramps/cli/test/user_test.py +++ b/gramps/cli/test/user_test.py @@ -151,5 +151,32 @@ class TestUser_quiet(unittest.TestCase): assert len(self.user._fileout.method_calls ) == 0, list(self.user._fileout.method_calls) +@unittest.skipUnless(MOCKING, "Requires unittest.mock to run") +class TestUser_progress(unittest.TestCase): + _progress_begin_step_end = \ + TestUser_quiet.test_progress_can_begin_step_end.__func__ + + def setUp(self): + self.user = user.User() + self.user._fileout = Mock(spec=sys.stderr) + # Collect baseline output from the old-style interface (begin/step/end) + self._progress_begin_step_end() + self.expected_output = list(self.user._fileout.method_calls) + self.user._fileout.reset_mock() + self.assertTrue( + len(self.user._fileout.method_calls) == 0, + list(self.user._fileout.method_calls)) + + def test(self): + with self.user.progress("Foo", "Bar", 0) as step: + for i in range(10): + step() + + def tearDown(self): + # Output using `with' differs from one with `progress_...' + self.assertEqual(self.expected_output, + list(self.user._fileout.method_calls)) + + if __name__ == "__main__": unittest.main() diff --git a/gramps/gen/user.py b/gramps/gen/user.py index a634983cf..1007c210c 100644 --- a/gramps/gen/user.py +++ b/gramps/gen/user.py @@ -21,6 +21,7 @@ # import sys +from contextlib import contextmanager """ The User class provides basic interaction with the user. @@ -83,6 +84,13 @@ class User(): Stop showing the progress indicator to the user. """ pass + + # Context-manager wrapper of the begin/step/end_progress above + @contextmanager + def progress(self, *args, **kwargs): + self.begin_progress(*args, **kwargs) + yield self.step_progress + self.end_progress() def prompt(self, title, message, accept_label, reject_label): """