DI-engine/ding/utils/tests/test_profiler_helper.py

from easydict import EasyDict
import pytest
import unittest
from unittest import mock
from unittest.mock import patch
import pathlib as pl
import os
import shutil

from ding.utils.profiler_helper import Profiler, register_profiler


@pytest.mark.unittest
class TestProfilerModule:

    def assertIsFile(self, path):
        if not pl.Path(path).resolve().is_file():
            raise AssertionError("File does not exist: %s" % str(path))

    def test(self):
        profiler = Profiler()

        def register_mock(write_profile, pr, folder_path):
            profiler.write_profile(pr, folder_path)

        def clean_up(dir):
            if os.path.exists(dir):
                shutil.rmtree(dir)

        dir = "./tmp_test/"
        clean_up(dir)

        with patch('ding.utils.profiler_helper.register_profiler', register_mock):
            profiler.profile(dir)
            file_path = os.path.join(dir, "profile_tottime.txt")
            self.assertIsFile(file_path)
            file_path = os.path.join(dir, "profile_cumtime.txt")
            self.assertIsFile(file_path)
            file_path = os.path.join(dir, "profile.prof")
            self.assertIsFile(file_path)

        clean_up(dir)