Setting query context has to be done with a separate method call

4 files changed, 88 insertions, 87 deletions

diff --git a/tests/test_api.py b/tests/test_api.py
index 4f248441..b2143548 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -539,10 +539,10 @@ class ApiTest(CoverageTest):
         filenames = self.get_measured_filenames(data)
         suite_filename = filenames['testsuite.py']
 
-        self.assertEqual(
-            [2, 8], data.lines(suite_filename, contexts=["multiply_six"]))
-        self.assertEqual(
-            [2, 5], data.lines(suite_filename, contexts=["multiply_zero"]))
+        data.set_query_context("multiply_six")
+        self.assertEqual([2, 8], data.lines(suite_filename))
+        data.set_query_context("multiply_zero")
+        self.assertEqual([2, 5], data.lines(suite_filename))
 
     def test_switch_context_with_static(self):
         # This test simulates a coverage-aware test runner,
@@ -579,10 +579,10 @@ class ApiTest(CoverageTest):
         filenames = self.get_measured_filenames(data)
         suite_filename = filenames['testsuite.py']
 
-        self.assertEqual(
-            [2, 8], data.lines(suite_filename, contexts=["mysuite|multiply_six"]))
-        self.assertEqual(
-            [2, 5], data.lines(suite_filename, contexts=["mysuite|multiply_zero"]))
+        data.set_query_context("mysuite|multiply_six")
+        self.assertEqual([2, 8], data.lines(suite_filename))
+        data.set_query_context("mysuite|multiply_zero")
+        self.assertEqual([2, 5], data.lines(suite_filename))
 
     def test_switch_context_unstarted(self):
         # Coverage must be started to switch context
diff --git a/tests/test_context.py b/tests/test_context.py
index eb8ee9c5..72739ad0 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -72,10 +72,14 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            self.assertEqual(combined.lines(fred, contexts=['red']), self.LINES)
-            self.assertEqual(combined.lines(fred, contexts=['blue']), [])
-            self.assertEqual(combined.lines(fblue, contexts=['red']), [])
-            self.assertEqual(combined.lines(fblue, contexts=['blue']), self.LINES)
+            def assert_combined_lines(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.lines(filename), lines)
+
+            assert_combined_lines(fred, 'red', self.LINES)
+            assert_combined_lines(fred, 'blue', [])
+            assert_combined_lines(fblue, 'red', [])
+            assert_combined_lines(fblue, 'blue', self.LINES)
 
     def test_combining_arc_contexts(self):
         red_data, blue_data = self.run_red_blue(branch=True)
@@ -92,15 +96,23 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            self.assertEqual(combined.lines(fred, contexts=['red']), self.LINES)
-            self.assertEqual(combined.lines(fred, contexts=['blue']), [])
-            self.assertEqual(combined.lines(fblue, contexts=['red']), [])
-            self.assertEqual(combined.lines(fblue, contexts=['blue']), self.LINES)
+            def assert_combined_lines(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.lines(filename), lines)
+
+            assert_combined_lines(fred, 'red', self.LINES)
+            assert_combined_lines(fred, 'blue', [])
+            assert_combined_lines(fblue, 'red', [])
+            assert_combined_lines(fblue, 'blue', self.LINES)
+
+            def assert_combined_arcs(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.arcs(filename), lines)
 
-            self.assertEqual(combined.arcs(fred, contexts=['red']), self.ARCS)
-            self.assertEqual(combined.arcs(fred, contexts=['blue']), [])
-            self.assertEqual(combined.arcs(fblue, contexts=['red']), [])
-            self.assertEqual(combined.arcs(fblue, contexts=['blue']), self.ARCS)
+            assert_combined_arcs(fred, 'red', self.ARCS)
+            assert_combined_arcs(fred, 'blue', [])
+            assert_combined_arcs(fblue, 'red', [])
+            assert_combined_arcs(fblue, 'blue', self.ARCS)
 
 
 class DynamicContextTest(CoverageTest):
@@ -145,12 +157,14 @@ class DynamicContextTest(CoverageTest):
         self.assertCountEqual(
             data.measured_contexts(),
             ["", "two_tests.test_one", "two_tests.test_two"])
-        self.assertCountEqual(data.lines(fname, [""]), self.OUTER_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["two_tests.test_one"]),
-            self.TEST_ONE_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["two_tests.test_two"]), self.TEST_TWO_LINES)
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertCountEqual(lines, data.lines(fname))
+
+        assert_context_lines("", self.OUTER_LINES)
+        assert_context_lines("two_tests.test_one", self.TEST_ONE_LINES)
+        assert_context_lines("two_tests.test_two", self.TEST_TWO_LINES)
 
     def test_static_and_dynamic(self):
         self.make_file("two_tests.py", self.SOURCE)
@@ -164,12 +178,14 @@ class DynamicContextTest(CoverageTest):
         self.assertCountEqual(
             data.measured_contexts(),
             ["stat", "stat|two_tests.test_one", "stat|two_tests.test_two"])
-        self.assertCountEqual(
-            data.lines(fname, ["stat"]), self.OUTER_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["stat|two_tests.test_one"]), self.TEST_ONE_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["stat|two_tests.test_two"]), self.TEST_TWO_LINES)
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertCountEqual(lines, data.lines(fname))
+
+        assert_context_lines("stat", self.OUTER_LINES)
+        assert_context_lines("stat|two_tests.test_one", self.TEST_ONE_LINES)
+        assert_context_lines("stat|two_tests.test_two", self.TEST_TWO_LINES)
 
 
 def get_qualname():
diff --git a/tests/test_data.py b/tests/test_data.py
index eb389e35..71e18bb9 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -201,8 +201,10 @@ class CoverageDataTest(DataTestHelpers, CoverageTest):
         covdata.set_context('test_a')
         covdata.add_lines(LINES_1)
         self.assertEqual(covdata.lines('a.py'), [1, 2])
-        self.assertEqual(covdata.lines('a.py', contexts=['test*']), [1, 2])
-        self.assertEqual(covdata.lines('a.py', contexts=['other*']), [])
+        covdata.set_query_contexts(['test*'])
+        self.assertEqual(covdata.lines('a.py'), [1, 2])
+        covdata.set_query_contexts(['other*'])
+        self.assertEqual(covdata.lines('a.py'), [])
 
     def test_contexts_by_lineno_with_lines(self):
         covdata = CoverageData()
@@ -250,11 +252,11 @@ class CoverageDataTest(DataTestHelpers, CoverageTest):
         covdata = CoverageData()
         covdata.set_context('test_x')
         covdata.add_arcs(ARCS_3)
-        self.assertEqual(
-            covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
-        self.assertEqual(covdata.arcs(
-            'x.py', contexts=['test*']), [(-1, 1), (1, 2), (2, 3), (3, -1)])
-        self.assertEqual(covdata.arcs('x.py', contexts=['other*']), [])
+        self.assertEqual(covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
+        covdata.set_query_contexts(['test*'])
+        self.assertEqual(covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
+        covdata.set_query_contexts(['other*'])
+        self.assertEqual(covdata.arcs('x.py'), [])
 
     def test_contexts_by_lineno_with_arcs(self):
         covdata = CoverageData()
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index eb51c1b9..22796dc3 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -990,18 +990,12 @@ class DynamicContextPluginTest(CoverageTest):
             ['', 'doctest:HTML_TAG', 'test:HTML_TAG', 'test:RENDERERS'],
             sorted(data.measured_contexts()),
         )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["test:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["test:RENDERERS"]),
-        )
+        data.set_query_context("doctest:HTML_TAG")
+        self.assertEqual([2], data.lines(filenames['rendering.py']))
+        data.set_query_context("test:HTML_TAG")
+        self.assertEqual([2], data.lines(filenames['rendering.py']))
+        data.set_query_context("test:RENDERERS")
+        self.assertEqual([2, 5, 8, 11], data.lines(filenames['rendering.py']))
 
     def test_static_context(self):
         self.make_plugin_capitalized_testnames('plugin_tests.py')
@@ -1041,22 +1035,21 @@ class DynamicContextPluginTest(CoverageTest):
         # labeled by plugin_tests.
         data = cov.get_data()
         filenames = self.get_measured_filenames(data)
-        self.assertEqual(
-            ['', 'doctest:HTML_TAG', 'testsuite.test_html_tag', 'testsuite.test_renderers'],
-            sorted(data.measured_contexts()),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["testsuite.test_html_tag"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["testsuite.test_renderers"]),
-        )
+        expected = [
+            '',
+            'doctest:HTML_TAG',
+            'testsuite.test_html_tag',
+            'testsuite.test_renderers',
+        ]
+        self.assertEqual(expected, sorted(data.measured_contexts()))
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertEqual(lines, data.lines(filenames['rendering.py']))
+
+        assert_context_lines("doctest:HTML_TAG", [2])
+        assert_context_lines("testsuite.test_html_tag", [2])
+        assert_context_lines("testsuite.test_renderers", [2, 5, 8, 11])
 
     def test_multiple_plugins(self):
         self.make_plugin_capitalized_testnames('plugin_tests.py')
@@ -1087,23 +1080,13 @@ class DynamicContextPluginTest(CoverageTest):
             'test:RENDERERS',
         ]
         self.assertEqual(expected, sorted(data.measured_contexts()))
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["test:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["test:RENDERERS"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5],
-            data.lines(filenames['rendering.py'], contexts=["renderer:paragraph"]),
-        )
-        self.assertEqual(
-            [2, 8],
-            data.lines(filenames['rendering.py'], contexts=["renderer:span"]),
-        )
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertEqual(lines, data.lines(filenames['rendering.py']))
+
+        assert_context_lines("test:HTML_TAG", [2])
+        assert_context_lines("test:RENDERERS", [2, 5, 8, 11])
+        assert_context_lines("doctest:HTML_TAG", [2])
+        assert_context_lines("renderer:paragraph", [2, 5])
+        assert_context_lines("renderer:span", [2, 8])