77import pytest
88
99from codeflash .code_utils .code_utils import (
10+ _get_site_packages_paths ,
1011 cleanup_paths ,
1112 exit_with_message ,
1213 file_name_from_test_module_name ,
@@ -263,6 +264,7 @@ def test_get_run_tmp_file_reuses_temp_directory() -> None:
263264def test_path_belongs_to_site_packages_with_site_package_path (monkeypatch : pytest .MonkeyPatch ) -> None :
264265 site_packages = [Path ("/usr/local/lib/python3.9/site-packages" ).resolve ()]
265266 monkeypatch .setattr (site , "getsitepackages" , lambda : site_packages )
267+ _get_site_packages_paths .cache_clear ()
266268
267269 file_path = Path ("/usr/local/lib/python3.9/site-packages/some_package" )
268270 assert path_belongs_to_site_packages (file_path ) is True
@@ -271,6 +273,7 @@ def test_path_belongs_to_site_packages_with_site_package_path(monkeypatch: pytes
271273def test_path_belongs_to_site_packages_with_non_site_package_path (monkeypatch : pytest .MonkeyPatch ) -> None :
272274 site_packages = [Path ("/usr/local/lib/python3.9/site-packages" )]
273275 monkeypatch .setattr (site , "getsitepackages" , lambda : site_packages )
276+ _get_site_packages_paths .cache_clear ()
274277
275278 file_path = Path ("/usr/local/lib/python3.9/other_directory/some_package" )
276279 assert path_belongs_to_site_packages (file_path ) is False
@@ -279,6 +282,7 @@ def test_path_belongs_to_site_packages_with_non_site_package_path(monkeypatch: p
279282def test_path_belongs_to_site_packages_with_relative_path (monkeypatch : pytest .MonkeyPatch ) -> None :
280283 site_packages = [Path ("/usr/local/lib/python3.9/site-packages" )]
281284 monkeypatch .setattr (site , "getsitepackages" , lambda : site_packages )
285+ _get_site_packages_paths .cache_clear ()
282286
283287 file_path = Path ("some_package" )
284288 assert path_belongs_to_site_packages (file_path ) is False
@@ -298,6 +302,7 @@ def test_path_belongs_to_site_packages_with_symlinked_site_packages(
298302 package_file .write_text ("# package file" )
299303
300304 monkeypatch .setattr (site , "getsitepackages" , lambda : [str (symlinked_site_packages )])
305+ _get_site_packages_paths .cache_clear ()
301306
302307 assert path_belongs_to_site_packages (package_file ) is True
303308
@@ -321,6 +326,7 @@ def test_path_belongs_to_site_packages_with_complex_symlinks(monkeypatch: pytest
321326
322327 site_packages_via_links = link2 / "lib" / "python3.9" / "site-packages"
323328 monkeypatch .setattr (site , "getsitepackages" , lambda : [str (site_packages_via_links )])
329+ _get_site_packages_paths .cache_clear ()
324330
325331 assert path_belongs_to_site_packages (package_file ) is True
326332
@@ -341,6 +347,7 @@ def test_path_belongs_to_site_packages_resolved_paths_normalization(
341347
342348 complex_site_packages_path = tmp_path / "lib" / "python3.9" / "other" / ".." / "site-packages" / "."
343349 monkeypatch .setattr (site , "getsitepackages" , lambda : [str (complex_site_packages_path )])
350+ _get_site_packages_paths .cache_clear ()
344351
345352 assert path_belongs_to_site_packages (package_file ) is True
346353
@@ -380,18 +387,16 @@ def my_function():
380387
381388
382389@pytest .fixture
383- def mock_code_context ():
390+ def mock_code_context () -> MagicMock :
384391 """Mock CodeOptimizationContext for testing extract_dependent_function."""
385- from unittest .mock import MagicMock
386-
387392 from codeflash .models .models import CodeOptimizationContext
388393
389394 context = MagicMock (spec = CodeOptimizationContext )
390395 context .preexisting_objects = []
391396 return context
392397
393398
394- def test_extract_dependent_function_sync_and_async (mock_code_context ) :
399+ def test_extract_dependent_function_sync_and_async (mock_code_context : MagicMock ) -> None :
395400 """Test extract_dependent_function with both sync and async functions."""
396401 # Test sync function extraction
397402 mock_code_context .testgen_context = CodeStringsMarkdown .parse_markdown_code ("""```python:file.py
@@ -417,7 +422,7 @@ async def async_helper_function():
417422 assert extract_dependent_function ("main_function" , mock_code_context ) == "async_helper_function"
418423
419424
420- def test_extract_dependent_function_edge_cases (mock_code_context ) :
425+ def test_extract_dependent_function_edge_cases (mock_code_context : MagicMock ) -> None :
421426 """Test extract_dependent_function edge cases."""
422427 # No dependent functions
423428 mock_code_context .testgen_context = CodeStringsMarkdown .parse_markdown_code ("""```python:file.py
@@ -441,7 +446,7 @@ async def helper2():
441446 assert extract_dependent_function ("main_function" , mock_code_context ) is False
442447
443448
444- def test_extract_dependent_function_mixed_scenarios (mock_code_context ) :
449+ def test_extract_dependent_function_mixed_scenarios (mock_code_context : MagicMock ) -> None :
445450 """Test extract_dependent_function with mixed sync/async scenarios."""
446451 # Async main with sync helper
447452 mock_code_context .testgen_context = CodeStringsMarkdown .parse_markdown_code ("""```python:file.py
0 commit comments