diff --git a/tests/conftest.py b/tests/conftest.py index 5e97102bdf..684dc7be0c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -70,6 +70,12 @@ def pytest_addoption(parser): "--benchmark", action="store_true", help="Collect benchmarking data for tests" ) + parser.addoption( + "--performance-report", + action="store_true", + help="Collect performance report for tests", + ) + parser.addoption( "--memray", action="store", @@ -910,3 +916,33 @@ def _memray_profile(client): s3.put(archive, s3_performance_url) yield _memray_profile + + +@pytest.fixture +def performance_report( + pytestconfig, + s3, + s3_performance_url, + s3_storage_options, + test_run_benchmark, + tmp_path, +): + if not test_run_benchmark: + yield + else: + if not pytestconfig.getoption("--performance-report"): + yield contextlib.nullcontext + else: + + @contextlib.contextmanager + def _performance_report(): + try: + filename = f"{s3_performance_url}/performance_report.html.gz" + with distributed.performance_report( + filename=filename, storage_options=s3_storage_options + ): + yield + finally: + test_run_benchmark.performance_report_url = filename + + yield _performance_report diff --git a/tests/geospatial/conftest.py b/tests/geospatial/conftest.py index 6c2ad73a2f..9af50f6624 100644 --- a/tests/geospatial/conftest.py +++ b/tests/geospatial/conftest.py @@ -27,7 +27,9 @@ def cluster_name(request, scale): @pytest.fixture() -def client_factory(cluster_name, github_cluster_tags, benchmark_all, memray_profile): +def client_factory( + cluster_name, github_cluster_tags, benchmark_all, memray_profile, performance_report +): import contextlib @contextlib.contextmanager @@ -43,7 +45,9 @@ def _(n_workers, env=None, **cluster_kwargs): with cluster.get_client() as client: # FIXME https://github.com/coiled/platform/issues/103 client.wait_for_workers(n_workers) - with memray_profile(client), benchmark_all(client): + with performance_report(), memray_profile(client), benchmark_all( + client + ): yield client return _