Expose switch_context in coverage API (#782)

* Expose switch_context in public API

* Test switch_context

* Helper method to get full paths to measured files

* Get correct file paths on all OS

* Note version that introduced this method

2 files changed, 127 insertions, 0 deletions

diff --git a/coverage/control.py b/coverage/control.py
index 9bcab844..ace3d8bf 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -483,6 +483,22 @@ class Coverage(object):
         self._data.erase(parallel=self.config.parallel)
         self._data = None
 
+    def switch_context(self, new_context):
+        """Switch to a new dynamic context.
+
+        `new_context` is a string to use as the context label
+        for collected data.
+
+        Coverage collection must be started already.
+
+        .. versionadded:: 5.0
+        """
+        if not self._started:
+            raise CoverageException(                    # pragma: only jython
+                "Cannot switch context, coverage is not started"
+                )
+        self._collector.switch_context(new_context)
+
     def clear_exclude(self, which='exclude'):
         """Clear the exclude list."""
         self._init()
diff --git a/tests/test_api.py b/tests/test_api.py
index e838e6b7..a1c98f04 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -505,6 +505,117 @@ class ApiTest(CoverageTest):
             b.py        1      0   100%
             """))
 
+    def make_testsuite(self):
+        """Create a simple file representing a method with two tests.
+
+        Returns absolute path to the file.
+        """
+        self.make_file("testsuite.py", """\
+            def timestwo(x):
+                return x*2
+
+            def test_multiply_zero():
+                assert timestwo(0) == 0
+
+            def test_multiply_six():
+                assert timestwo(6) == 12
+            """)
+
+    def test_switch_context_testrunner(self):
+        # This test simulates a coverage-aware test runner,
+        # measuring labeled coverage via public API
+        self.make_testsuite()
+
+        # Test runner starts
+        cov = coverage.Coverage()
+        cov.start()
+
+        # Imports the test suite
+        suite = import_local_file("testsuite")
+
+        # Measures test case 1
+        cov.switch_context('multiply_zero')
+        suite.test_multiply_zero()
+
+        # Measures test case 2
+        cov.switch_context('multiply_six')
+        suite.test_multiply_six()
+
+        # Runner finishes
+        cov.save()
+        cov.stop()
+
+        # Labeled data is collected
+        data = cov.get_data()
+        self.assertEqual(
+            sorted(data.measured_contexts()),
+            [u'', u'multiply_six', u'multiply_zero'])
+
+        filenames = self.get_measured_filenames(data)
+        suite_filename = filenames['testsuite.py']
+
+        self.assertEqual(
+            data.lines(suite_filename, context="multiply_six"),
+            [2, 8])
+        self.assertEqual(
+            data.lines(suite_filename, context="multiply_zero"),
+            [2, 5])
+
+    def test_switch_context_with_static(self):
+        # This test simulates a coverage-aware test runner,
+        # measuring labeled coverage via public API,
+        # with static label prefix.
+        self.make_testsuite()
+
+        # Test runner starts
+        cov = coverage.Coverage(context="mysuite")
+        cov.start()
+
+        # Imports the test suite
+        suite = import_local_file("testsuite")
+
+        # Measures test case 1
+        cov.switch_context('multiply_zero')
+        suite.test_multiply_zero()
+
+        # Measures test case 2
+        cov.switch_context('multiply_six')
+        suite.test_multiply_six()
+
+        # Runner finishes
+        cov.save()
+        cov.stop()
+
+        # Labeled data is collected
+        data = cov.get_data()
+        self.assertEqual(
+            sorted(data.measured_contexts()),
+            [u'mysuite', u'mysuite:multiply_six', u'mysuite:multiply_zero'])
+
+        filenames = self.get_measured_filenames(data)
+        suite_filename = filenames['testsuite.py']
+
+        self.assertEqual(
+            data.lines(suite_filename, context="mysuite:multiply_six"),
+            [2, 8])
+        self.assertEqual(
+            data.lines(suite_filename, context="mysuite:multiply_zero"),
+            [2, 5])
+
+    def test_switch_context_unstarted(self):
+        # Coverage must be started to switch context
+
+        cov = coverage.Coverage()
+        with self.assertRaises(CoverageException):
+            cov.switch_context("test1")
+
+        cov.start()
+        cov.switch_context("test2")
+
+        cov.stop()
+        with self.assertRaises(CoverageException):
+            cov.switch_context("test3")
+
 
 class NamespaceModuleTest(UsingModulesMixin, CoverageTest):
     """Test PEP-420 namespace modules."""