Multiprocessing needs to communicate the rcfile down to the subprocesses

3 files changed, 63 insertions, 20 deletions

diff --git a/coverage/control.py b/coverage/control.py
index e0e2e6f..fed5ab4 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -133,6 +133,7 @@ class Coverage(object):
             # True, so make it so.
             if config_file == ".coveragerc":
                 config_file = True
+            self.config_file = config_file
             specified_file = (config_file is not True)
             if not specified_file:
                 config_file = ".coveragerc"
@@ -251,7 +252,7 @@ class Coverage(object):
 
         concurrency = self.config.concurrency or []
         if "multiprocessing" in concurrency:
-            patch_multiprocessing()
+            patch_multiprocessing(rcfile=self.config_file)
             #concurrency = None
             # Multi-processing uses parallel for the subprocesses, so also use
             # it for the main process.
diff --git a/coverage/multiproc.py b/coverage/multiproc.py
index d0bdf00..f9341ef 100644
--- a/coverage/multiproc.py
+++ b/coverage/multiproc.py
@@ -1,7 +1,7 @@
 # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
 # For details: https://bitbucket.org/ned/coveragepy/src/default/NOTICE.txt
 
-"""Monkey-patching to make coverage.py work right in some cases."""
+"""Monkey-patching to add multiprocessing support for coverage.py"""
 
 import multiprocessing
 import multiprocessing.process
@@ -9,22 +9,23 @@ import sys
 
 # An attribute that will be set on modules to indicate that they have been
 # monkey-patched.
-PATCHED_MARKER = "_coverage$patched"
+PATCHED_MARKER = "_coverage$rcfile"
 
 if sys.version_info >= (3, 4):
-    klass = multiprocessing.process.BaseProcess
+    OriginalProcess = multiprocessing.process.BaseProcess
 else:
-    klass = multiprocessing.Process
+    OriginalProcess = multiprocessing.Process
 
-original_bootstrap = klass._bootstrap
+original_bootstrap = OriginalProcess._bootstrap
 
 
-class ProcessWithCoverage(klass):
+class ProcessWithCoverage(OriginalProcess):
     """A replacement for multiprocess.Process that starts coverage."""
     def _bootstrap(self):
         """Wrapper around _bootstrap to start coverage."""
         from coverage import Coverage
-        cov = Coverage(data_suffix=True)
+        rcfile = getattr(multiprocessing, PATCHED_MARKER)
+        cov = Coverage(data_suffix=True, config_file=rcfile)
         cov.start()
         try:
             return original_bootstrap(self)
@@ -35,25 +36,30 @@ class ProcessWithCoverage(klass):
 
 class Stowaway(object):
     """An object to pickle, so when it is unpickled, it can apply the monkey-patch."""
+    def __init__(self, rcfile):
+        self.rcfile = rcfile
+
     def __getstate__(self):
-        return {}
+        return {'rcfile': self.rcfile}
 
-    def __setstate__(self, state_unused):
-        patch_multiprocessing()
+    def __setstate__(self, state):
+        patch_multiprocessing(state['rcfile'])
 
 
-def patch_multiprocessing():
+def patch_multiprocessing(rcfile):
     """Monkey-patch the multiprocessing module.
 
     This enables coverage measurement of processes started by multiprocessing.
-    This is wildly experimental!
+    This involves aggressive monkey-patching.
+
+    `rcfile` is the path to the rcfile being used.
 
     """
     if hasattr(multiprocessing, PATCHED_MARKER):
         return
 
     if sys.version_info >= (3, 4):
-        klass._bootstrap = ProcessWithCoverage._bootstrap
+        OriginalProcess._bootstrap = ProcessWithCoverage._bootstrap
     else:
         multiprocessing.Process = ProcessWithCoverage
 
@@ -72,9 +78,9 @@ def patch_multiprocessing():
         def get_preparation_data_with_stowaway(name):
             """Get the original preparation data, and also insert our stowaway."""
             d = original_get_preparation_data(name)
-            d['stowaway'] = Stowaway()
+            d['stowaway'] = Stowaway(rcfile)
             return d
 
         spawn.get_preparation_data = get_preparation_data_with_stowaway
 
-    setattr(multiprocessing, PATCHED_MARKER, True)
+    setattr(multiprocessing, PATCHED_MARKER, rcfile)
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 3603443..d438684 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -324,7 +324,7 @@ MULTI_CODE = """
         ret = work(*args)
         return os.getpid(), ret
 
-    if __name__ == "__main__":
+    if __name__ == "__main__":      # pragma: no branch
         # This if is on a single line so we can get 100% coverage
         # even if we have no arguments.
         if len(sys.argv) > 1: multiprocessing.set_start_method(sys.argv[1])
@@ -372,12 +372,12 @@ class MultiprocessingTest(CoverageTest):
             else:
                 self.assertEqual(out.rstrip(), expected_out)
 
-                self.run_command("coverage combine")
+                out = self.run_command("coverage combine")
+                self.assertEqual(out, "")
                 out = self.run_command("coverage report -m")
 
                 last_line = self.squeezed_lines(out)[-1]
-                expected_report = "multi.py {lines} 0 100%".format(lines=line_count(code))
-                self.assertEqual(last_line, expected_report)
+                self.assertRegex(last_line, r"multi.py \d+ 0 100%")
 
     def test_multiprocessing(self):
         nprocs = 3
@@ -398,3 +398,39 @@ class MultiprocessingTest(CoverageTest):
         self.try_multiprocessing_code(
             code, expected_out, eventlet, concurrency="multiprocessing,eventlet"
         )
+
+    def try_multiprocessing_code_with_branching(self, code, expected_out):
+        """Run code using multiprocessing, it should produce `expected_out`."""
+        self.make_file("multi.py", code)
+        self.make_file("multi.rc", """\
+            [run]
+            concurrency = multiprocessing
+            branch = True
+            """)
+
+        if env.PYVERSION >= (3, 4):
+            start_methods = ['fork', 'spawn']
+        else:
+            start_methods = ['']
+
+        for start_method in start_methods:
+            if start_method and start_method not in multiprocessing.get_all_start_methods():
+                continue
+
+            out = self.run_command("coverage run --rcfile=multi.rc multi.py %s" % (start_method,))
+            self.assertEqual(out.rstrip(), expected_out)
+
+            out = self.run_command("coverage combine")
+            self.assertEqual(out, "")
+            out = self.run_command("coverage report -m")
+
+            last_line = self.squeezed_lines(out)[-1]
+            self.assertRegex(last_line, r"multi.py \d+ 0 \d+ 0 100%")
+
+    def test_multiprocessing_with_branching(self):
+        nprocs = 3
+        upto = 30
+        code = (SQUARE_OR_CUBE_WORK + MULTI_CODE).format(NPROCS=nprocs, UPTO=upto)
+        total = sum(x*x if x%2 else x*x*x for x in range(upto))
+        expected_out = "{nprocs} pids, total = {total}".format(nprocs=nprocs, total=total)
+        self.try_multiprocessing_code_with_branching(code, expected_out)