add tests

Author: KRRT7

codeflash/verification/pytest_plugin.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-import importlib
+import contextlib
 import inspect
 
 # System Imports
 import logging
 import os
 import re
+import sys
 import time
 import warnings
 from typing import TYPE_CHECKING, Any, Callable
@@ -136,9 +137,6 @@ def pytest_runtestloop(self, session: Session) -> bool:
                     item._nodeid = self._set_nodeid(item._nodeid, count)
 
                 next_item: pytest.Item = session.items[index + 1] if index + 1 < len(session.items) else None
-
-                self._clear_lru_caches(next_item)
-
                 item.config.hook.pytest_runtest_protocol(item=item, nextitem=next_item)
                 if session.shouldfail:
                     raise session.Failed(session.shouldfail)
@@ -149,6 +147,58 @@ def pytest_runtestloop(self, session: Session) -> bool:
             time.sleep(self._get_delay_time(session))
         return True
 
+    def _clear_lru_caches(self, item: pytest.Item) -> None:
+        processed_functions: set[Callable] = set()
+        protected_modules = {
+            "gc",
+            "inspect",
+            "os",
+            "sys",
+            "time",
+            "functools",
+            "pathlib",
+            "typing",
+            "dill",
+            "pytest",
+            "importlib",
+        }
+
+        def _clear_cache_for_object(obj: Any) -> None:  # noqa: ANN401
+            if obj in processed_functions:
+                return
+            processed_functions.add(obj)
+
+            if hasattr(obj, "__wrapped__"):
+                module_name = obj.__wrapped__.__module__
+            else:
+                try:
+                    obj_module = inspect.getmodule(obj)
+                    module_name = obj_module.__name__.split(".")[0] if obj_module is not None else None
+                except Exception:  # noqa: BLE001
+                    module_name = None
+
+            if module_name in protected_modules:
+                return
+
+            if hasattr(obj, "cache_clear") and callable(obj.cache_clear):
+                with contextlib.suppress(Exception):
+                    obj.cache_clear()
+
+        _clear_cache_for_object(item.function)  # type: ignore[attr-defined]
+
+        try:
+            if hasattr(item.function, "__module__"):  # type: ignore[attr-defined]
+                module_name = item.function.__module__  # type: ignore[attr-defined]
+                try:
+                    module = sys.modules.get(module_name)
+                    if module:
+                        for _, obj in inspect.getmembers(module):
+                            if callable(obj):
+                                _clear_cache_for_object(obj)
+                except Exception:  # noqa: BLE001, S110
+                    pass
+        except Exception:  # noqa: BLE001, S110
+            pass
     def _set_nodeid(self, nodeid: str, count: int) -> str:
         """Set loop count when using duration.
 
@@ -194,72 +244,6 @@ def _timed_out(self, session: Session, start_time: float, count: int) -> bool:
             and time.time() - start_time > self._get_total_time(session)
         )
 
-    def _clear_lru_caches(self, item: pytest.Item) -> None:
-        processed_functions: set[Callable] = set()
-        cleared_caches = 0
-        protected_modules = {
-            "gc",
-            "inspect",
-            "os",
-            "sys",
-            "time",
-            "functools",
-            "pathlib",
-            "typing",
-            "dill",
-            "pytest",
-            "importlib",
-        }
-
-        def _clear_cache_for_object(obj: Any) -> None:
-            if obj in processed_functions:
-                return
-            processed_functions.add(obj)
-
-            try:
-                obj_module = inspect.getmodule(obj)
-                module_name = obj_module.__name__.split(".")[0] if obj_module is not None else None
-            except AttributeError:
-                module_name = None
-
-            if module_name in protected_modules:
-                return
-
-            if hasattr(obj, "cache_clear") and callable(obj.cache_clear):
-                try:
-                    obj.cache_clear()
-                    nonlocal cleared_caches
-                    cleared_caches += 1
-                    print(f"Cleared cache for: ^{obj.__name__}^")
-                except Exception:  # noqa: BLE001, S110
-                    pass
-
-            if callable(obj):
-                try:
-                    for name in obj.__globals__:
-                        nested_obj = obj.__globals__[name]
-                        if callable(nested_obj):
-                            _clear_cache_for_object(nested_obj)
-                except Exception:  # noqa: BLE001, S110
-                    pass
-
-        if hasattr(item, "function") and callable(item.function):
-            module_name = item.function.__module__
-            full_file_path = inspect.getfile(item.function)
-            try:
-                module = importlib.import_module(module_name)
-
-                for _, obj in inspect.getmembers(module):
-                    if callable(obj):
-                        _clear_cache_for_object(obj)
-
-                    if isinstance(obj, type):
-                        for _, class_method in inspect.getmembers(obj, predicate=inspect.isfunction):
-                            _clear_cache_for_object(class_method)
-
-            except Exception:  # noqa: BLE001, S110
-                pass
-
     @pytest.fixture
     def __pytest_loop_step_number(self, request: pytest.FixtureRequest) -> int:
         """Set step number for loop.

tests/_sometime_test_lru_cache_clear.py

@@ -0,0 +1,87 @@
+import types
+from typing import NoReturn
+
+import pytest
+from _pytest.config import Config
+
+from codeflash.verification.pytest_plugin import PyTest_Loops
+
+
+@pytest.fixture
+def pytest_loops_instance(pytestconfig: Config) -> PyTest_Loops:
+    return PyTest_Loops(pytestconfig)
+
+
+@pytest.fixture
+def mock_item() -> type:
+    class MockItem:
+        def __init__(self, function: types.FunctionType) -> None:
+            self.function = function
+
+    return MockItem
+
+
+def create_mock_module(module_name: str, source_code: str) -> types.ModuleType:
+    module = types.ModuleType(module_name)
+    exec(source_code, module.__dict__)  # noqa: S102
+    return module
+
+
+def test_clear_lru_caches_function(pytest_loops_instance: PyTest_Loops, mock_item: type) -> None:
+    source_code = """
+import functools
+
+@functools.lru_cache(maxsize=None)
+def my_func(x):
+    return x * 2
+
+my_func(10)  # miss the cache
+my_func(10)  # hit the cache
+"""
+    mock_module = create_mock_module("test_module_func", source_code)
+    item = mock_item(mock_module.my_func)
+    pytest_loops_instance._clear_lru_caches(item)  # noqa: SLF001
+    assert mock_module.my_func.cache_info().hits == 0
+    assert mock_module.my_func.cache_info().misses == 0
+    assert mock_module.my_func.cache_info().currsize == 0
+
+
+def test_clear_lru_caches_class_method(pytest_loops_instance: PyTest_Loops, mock_item: type) -> None:
+    source_code = """
+import functools
+
+class MyClass:
+    @functools.lru_cache(maxsize=None)
+    def my_method(self, x):
+        return x * 3
+
+obj = MyClass()
+obj.my_method(5)  # Pre-populate the cache
+obj.my_method(5)  # Hit the cache
+# """
+    mock_module = create_mock_module("test_module_class", source_code)
+    item = mock_item(mock_module.MyClass.my_method)
+    pytest_loops_instance._clear_lru_caches(item)  # noqa: SLF001
+    assert mock_module.MyClass.my_method.cache_info().hits == 0
+    assert mock_module.MyClass.my_method.cache_info().misses == 0
+    assert mock_module.MyClass.my_method.cache_info().currsize == 0
+
+
+def test_clear_lru_caches_exception_handling(pytest_loops_instance: PyTest_Loops, mock_item: type) -> None:
+    """Test that exceptions during clearing are handled."""
+
+    class BrokenCache:
+        def cache_clear(self) -> NoReturn:
+            msg = "Cache clearing failed!"
+            raise ValueError(msg)
+
+    item = mock_item(BrokenCache())
+    pytest_loops_instance._clear_lru_caches(item)  # noqa: SLF001
+
+
+def test_clear_lru_caches_no_cache(pytest_loops_instance: PyTest_Loops, mock_item: type) -> None:
+    def no_cache_func(x: int) -> int:
+        return x
+
+    item = mock_item(no_cache_func)
+    pytest_loops_instance._clear_lru_caches(item)  # noqa: SLF001