add task_instrumentation fixture (#13176)

* add task_instrumentation fixture

* fixup

* Update conftest.py

chia/util/task_timing.py

@@ -257,11 +257,10 @@ def fontcolor(pct: float) -> str:
         return "black"
 
 
-def stop_task_instrumentation() -> None:
+def stop_task_instrumentation(target_dir: str = f"task-profile-{os.getpid()}") -> None:
     sys.setprofile(None)
     global g_function_infos
 
-    target_dir = f"task-profile-{os.getpid()}"
     try:
         os.mkdir(target_dir)
     except Exception:
@@ -335,12 +334,12 @@ def stop_task_instrumentation() -> None:
             f.write("}\n")
 
 
-if __name__ == "__main__":
+def main(args: List[str]) -> int:
     import glob
     import pathlib
     import subprocess
 
-    profile_dir = pathlib.Path(sys.argv[1])
+    profile_dir = pathlib.Path(args[0])
     queue: List[subprocess.Popen[bytes]] = []
     for file in glob.glob(str(profile_dir / "*.dot")):
         print(file)
@@ -357,3 +356,9 @@ if __name__ == "__main__":
     while len(queue) > 0:
         oldest = queue.pop(0)
         oldest.wait()
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(args=sys.argv[1:]))

tests/conftest.py

@@ -1,10 +1,13 @@
 # flake8: noqa E402 # See imports after multiprocessing.set_start_method
 import aiohttp
+import datetime
 import multiprocessing
 import os
 import sysconfig
-from typing import Any, AsyncIterator, Dict, List, Tuple, Union
+from typing import Any, AsyncIterator, Dict, Iterator, List, Tuple, Union
 
+# TODO: update after resolution in https://github.com/pytest-dev/pytest/issues/7469
+from _pytest.fixtures import SubRequest
 import pytest
 import pytest_asyncio
 import tempfile
@@ -27,6 +30,11 @@ from chia.simulator.simulator_protocol import FarmNewBlockProtocol
 from chia.types.peer_info import PeerInfo
 from chia.util.config import create_default_chia_config, lock_and_load_config
 from chia.util.ints import uint16
+from chia.util.task_timing import (
+    main as task_instrumentation_main,
+    start_task_instrumentation,
+    stop_task_instrumentation,
+)
 from chia.wallet.wallet import Wallet
 from tests.core.data_layer.util import ChiaRoot
 from tests.core.node_height import node_height_at_least
@@ -54,6 +62,27 @@ from tests.util.keyring import TempKeyring
 from tests.setup_nodes import setup_farmer_multi_harvester
 
 
+@pytest.fixture(name="node_name_for_file")
+def node_name_for_file_fixture(request: SubRequest) -> str:
+    # TODO: handle other characters banned on windows
+    return request.node.name.replace(os.sep, "_")
+
+
+@pytest.fixture(name="test_time_for_file")
+def test_time_for_file_fixture(request: SubRequest) -> str:
+    return datetime.datetime.now().isoformat().replace(":", "_")
+
+
+@pytest.fixture(name="task_instrumentation")
+def task_instrumentation_fixture(node_name_for_file: str, test_time_for_file: str) -> Iterator[None]:
+    target_directory = f"task-profile-{node_name_for_file}-{test_time_for_file}"
+
+    start_task_instrumentation()
+    yield
+    stop_task_instrumentation(target_dir=target_directory)
+    task_instrumentation_main(args=[target_directory])
+
+
 @pytest.fixture(scope="session")
 def get_keychain():
     with TempKeyring() as keychain: