unit test

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

src/compressed_tensors/utils/helpers.py

@@ -372,6 +372,7 @@ def patch_attr(base: object, attr: str, value: Any):
 @contextlib.contextmanager
 def patch_attrs(bases: list[object], attr: str, values: list[Any]):
     """
+    Same as `patch_attr` but for a list of objects to patch
     Patch attribute for a list of objects with list of values.
     Original values are restored upon exit
 
@@ -383,9 +384,11 @@ def patch_attrs(bases: list[object], attr: str, values: list[Any]):
     Usage:
     >>> from types import SimpleNamespace
     >>> obj = SimpleNamespace()
-    >>> with patch_attr(obj, "attribute", "value"):
-    ...     assert obj.attribute == "value"
-    >>> assert not hasattr(obj, "attribute")
+    >>> with patch_attr([obj1, obj2], "attribute", ["value1", "value2"]):
+    ...     assert obj1.attribute == "value1"
+    ...     assert obj2.attribute == "value2"
+    >>> assert not hasattr(obj1, "attribute")
+    >>> assert not hasattr(obj2, "attribute")
     """
     _sentinel = object()
     original_values = [getattr(base, attr, _sentinel) for base in bases]

tests/test_utils/test_helpers.py

@@ -21,6 +21,7 @@
     ParameterizedDefaultDict,
     load_compressed,
     patch_attr,
+    patch_attrs,
     save_compressed,
     save_compressed_model,
 )
@@ -176,6 +177,23 @@ def test_patch_attr():
     assert not hasattr(obj, "attribute")
 
 
+def test_patch_attrs():
+    num_objs = 4
+    objs = [SimpleNamespace() for _ in range(num_objs)]
+    for idx, obj in enumerate(objs):
+        if idx % 2 == 0:
+            obj.attribute = f"original_{idx}"
+    with patch_attrs(objs, "attribute", [f"patched_{idx}" for idx in range(num_objs)]):
+        for idx, obj in enumerate(objs):
+            assert obj.attribute == f"patched_{idx}"
+            obj.attribute = "modified"
+    for idx, obj in enumerate(objs):
+        if idx % 2 == 0:
+            assert obj.attribute == f"original_{idx}"
+        else:
+            assert not hasattr(obj, "attribute")
+
+
 def test_parameterized_default_dict():
     def add_one(value):
         return value + 1