diff --git a/python/.coveragerc b/python/.coveragerc index 111f791b04ba..cba7f5481af8 100644 --- a/python/.coveragerc +++ b/python/.coveragerc @@ -35,3 +35,8 @@ omit = */pyspark/streaming/tests/* */pyspark/tests/* */pyspark/testing/tests/* + +[paths] +source = + pyspark/ + lib/pyspark.zip/pyspark diff --git a/python/test_coverage/sitecustomize.py b/python/test_coverage/sitecustomize.py index 1f31860d2ff4..6bda12db883d 100644 --- a/python/test_coverage/sitecustomize.py +++ b/python/test_coverage/sitecustomize.py @@ -21,6 +21,32 @@ # variable is set or not. If set, it starts to run the coverage. try: import coverage - coverage.process_startup() + cov = coverage.process_startup() + if cov: + import os + + def patch_worker(): + # If it's a worker forked from the daemon, we need to patch it to save + # the coverage data. Otherwise the worker will be killed by a signal and + # the coverage data will not be saved. + import sys + frame = sys._getframe(1) + if ( + frame.f_code.co_name == "manager" and + "daemon.py" in frame.f_code.co_filename and + "worker" in frame.f_globals + ): + + def save_when_exit(func): + def wrapper(*args, **kwargs): + result = func(*args, **kwargs) + cov.save() + return result + return wrapper + + frame.f_globals["worker"] = save_when_exit(frame.f_globals["worker"]) + + os.register_at_fork(after_in_child=patch_worker) + except ImportError: pass