Merge pull request #1465 from codeflash-ai/brainiac

clear lru cache between runs

Author: KRRT7

codeflash/verification/pytest_plugin.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+import contextlib
+import inspect
+
 # System Imports
 import logging
 import os
 import re
+import sys
 import time
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Callable
 from unittest import TestCase
 
 # PyTest Imports
@@ -133,6 +137,9 @@ 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(item)
+
                 item.config.hook.pytest_runtest_protocol(item=item, nextitem=next_item)
                 if session.shouldfail:
                     raise session.Failed(session.shouldfail)
@@ -143,6 +150,59 @@ 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.
 
@@ -205,8 +265,7 @@ def __pytest_loop_step_number(self, request: pytest.FixtureRequest) -> int:
                     warnings.warn("Repeating unittest class tests not supported")
                 else:
                     raise UnexpectedError(
-                        "This call couldn't work with pytest-loops. "
-                        "Please consider raising an issue with your usage."
+                        "This call couldn't work with pytest-loops. Please consider raising an issue with your usage."
                     )
         return count
 
@@ -226,7 +285,7 @@ def pytest_generate_tests(self, metafunc: Metafunc) -> None:
             metafunc.fixturenames.append("__pytest_loop_step_number")
 
             def make_progress_id(i: int, n: int = count) -> str:
-                return f"{n}/{i+1}"
+                return f"{n}/{i + 1}"
 
             scope = metafunc.config.option.codeflash_loops_scope
             metafunc.parametrize(

tests/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