bpo-29235: Make cProfile.Profile a context manager (GH-6808)

1 files changed, 27 insertions, 0 deletions

diff --git a/Lib/test/test_cprofile.py b/Lib/test/test_cprofile.py
index 1430d22..2fd67ee 100644
--- a/Lib/test/test_cprofile.py
+++ b/Lib/test/test_cprofile.py
@@ -49,6 +49,33 @@ class CProfileTest(ProfileTest):
         # Test successful run
         assert_python_ok('-m', 'cProfile', '-m', 'timeit', '-n', '1')
 
+    def test_profile_enable_disable(self):
+        prof = self.profilerclass()
+        # Make sure we clean ourselves up if the test fails for some reason.
+        self.addCleanup(prof.disable)
+
+        prof.enable()
+        self.assertIs(sys.getprofile(), prof)
+
+        prof.disable()
+        self.assertIs(sys.getprofile(), None)
+
+    def test_profile_as_context_manager(self):
+        prof = self.profilerclass()
+        # Make sure we clean ourselves up if the test fails for some reason.
+        self.addCleanup(prof.disable)
+
+        with prof as __enter__return_value:
+            # profile.__enter__ should return itself.
+            self.assertIs(prof, __enter__return_value)
+
+            # profile should be set as the global profiler inside the
+            # with-block
+            self.assertIs(sys.getprofile(), prof)
+
+        # profile shouldn't be set once we leave the with-block.
+        self.assertIs(sys.getprofile(), None)
+
 
 def test_main():
     run_unittest(CProfileTest)