test_debug_dir.py

# Owner(s): ["module: dynamo"]
import shutil
import unittest

import torch
import torch._dynamo.test_case
import torch._dynamo.testing
from torch._dynamo.utils import DebugDir, get_debug_dir


class DebugDirTests(torch._dynamo.test_case.TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls._exit_stack.enter_context(
            unittest.mock.patch.object(
                torch._dynamo.config,
                "debug_dir_root",
                "/tmp/torch._dynamo_debug_dirs/",
            )
        )

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(torch._dynamo.config.debug_dir_root, ignore_errors=True)
        cls._exit_stack.close()

    def setUp(self):
        super().setUp()
        torch._dynamo.utils.debug_dir = DebugDir()

    def tearDown(self):
        torch._dynamo.utils.debug_dir = DebugDir()
        super().tearDown()

    def _setup(self):
        debug_dir = torch._dynamo.utils.debug_dir
        debug_dir.setup()
        self.assertIsNotNone(debug_dir.debug_path)
        self.assertEqual(debug_dir.num_setup_calls, 1)
        return debug_dir

    def test_setup(self):
        self._setup()

    def test_clear(self):
        debug_dir = self._setup()
        debug_dir.clear()
        self.assertIsNone(debug_dir.debug_path)
        self.assertEqual(debug_dir.num_setup_calls, 0)

    def test_multi_setup_single_clear(self):
        debug_dir = self._setup()
        prev = get_debug_dir()

        debug_dir.setup()
        self.assertEqual(prev, get_debug_dir())
        self.assertEqual(debug_dir.num_setup_calls, 2)

        debug_dir.clear()
        self.assertEqual(prev, get_debug_dir())
        self.assertEqual(debug_dir.num_setup_calls, 1)

    def test_multi_setup_multi_clear(self):
        debug_dir = self._setup()
        prev = get_debug_dir()

        debug_dir.setup()
        self.assertEqual(prev, get_debug_dir())
        self.assertEqual(debug_dir.num_setup_calls, 2)

        debug_dir.clear()
        self.assertEqual(prev, get_debug_dir())
        self.assertEqual(debug_dir.num_setup_calls, 1)

        debug_dir.clear()
        self.assertIsNone(debug_dir.debug_path)
        self.assertEqual(debug_dir.num_setup_calls, 0)

    def test_single_setup_single_clear(self):
        debug_dir = self._setup()
        debug_dir.clear()
        self.assertIsNone(debug_dir.debug_path)
        self.assertEqual(debug_dir.num_setup_calls, 0)

    def test_multi_get(self):
        self._setup()
        prev = get_debug_dir()
        next = get_debug_dir()
        self.assertEqual(prev, next)


if __name__ == "__main__":
    from torch._dynamo.test_case import run_tests

    run_tests()