From 241d2259d38fd67b289f65bf3a26a61231d33f51 Mon Sep 17 00:00:00 2001
From: "Colin L. Rice" <clr@fb.com>
Date: Tue, 19 Nov 2024 17:16:27 -0700
Subject: [PATCH] torch/config: fix mock behaviour (#140779)

Mock only replaces the value that was removed, if after deletion, it
does not see the attribute.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/140779
Approved by: https://github.com/ezyang
---
 test/test_utils_config_module.py |  9 +++++++++
 torch/utils/_config_module.py    | 17 +++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/test/test_utils_config_module.py b/test/test_utils_config_module.py
index add6af20b02e..50e391225b83 100644
--- a/test/test_utils_config_module.py
+++ b/test/test_utils_config_module.py
@@ -1,6 +1,7 @@
 # Owner(s): ["module: unknown"]
 import os
 import pickle
+from unittest.mock import patch
 
 
 os.environ["ENV_TRUE"] = "1"
@@ -320,6 +321,14 @@ torch.testing._internal.fake_config_module._save_config_ignore = ['e_ignored']""
         revert()
         self.assertTrue(config.e_bool)
 
+    def test_unittest_patch(self):
+        with patch("torch.testing._internal.fake_config_module.e_bool", False):
+            with patch("torch.testing._internal.fake_config_module.e_bool", False):
+                self.assertFalse(config.e_bool)
+            # unittest.mock has some very weird semantics around deletion of attributes when undoing patches
+            self.assertFalse(config.e_bool)
+        self.assertTrue(config.e_bool)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/torch/utils/_config_module.py b/torch/utils/_config_module.py
index 4dffd34a4e58..269ca4e59f82 100644
--- a/torch/utils/_config_module.py
+++ b/torch/utils/_config_module.py
@@ -213,6 +213,18 @@ class _ConfigEntry:
     # environment variables are read at install time
     env_value_force: Any = _UNSET_SENTINEL
     env_value_default: Any = _UNSET_SENTINEL
+    # Used to work arounds bad assumptions in unittest.mock.patch
+    # The code to blame is
+    # https://github.com/python/cpython/blob/94a7a4e22fb8f567090514785c69e65298acca42/Lib/unittest/mock.py#L1637
+    # Essentially, mock.patch requires, that if __dict__ isn't accessible
+    # (which it isn't), that after delattr is called on the object, the
+    # object must throw when hasattr is called. Otherwise, it doesn't call
+    # setattr again.
+    # Technically we'll have an intermediate state of hiding the config while
+    # mock.patch is unpatching itself, but it calls setattr after the delete
+    # call so the final state is correct. It's just very unintuitive.
+    # upstream bug - python/cpython#126886
+    hide: bool = False
 
     def __init__(self, config: Config):
         self.default = config.default
@@ -253,11 +265,15 @@ class ConfigModule(ModuleType):
         else:
             self._config[name].user_override = value
             self._is_dirty = True
+            self._config[name].hide = False
 
     def __getattr__(self, name: str) -> Any:
         try:
             config = self._config[name]
 
+            if config.hide:
+                raise AttributeError(f"{self.__name__}.{name} does not exist")
+
             if config.env_value_force is not _UNSET_SENTINEL:
                 return config.env_value_force
 
@@ -288,6 +304,7 @@ class ConfigModule(ModuleType):
         # must support delete because unittest.mock.patch deletes
         # then recreate things
         self._config[name].user_override = _UNSET_SENTINEL
+        self._config[name].hide = True
 
     def _is_default(self, name: str) -> bool:
         return self._config[name].user_override is _UNSET_SENTINEL