fix: include module-level variables and second-degree helpers in testgen helper FQNs

[KRRT7]

Summary

• Track module-level variable definitions ("statement" type) in Jedi helper resolution alongside functions and classes
• Add testgen_helper_fqns field to CodeOptimizationContext containing FQNs from both first-degree and second-degree helpers
• Use the complete FQN list for helper_function_names sent to the AI test generator

This fixes generated tests importing symbols (e.g., _EXTENSION_REGISTRY) from the wrong module, which caused ImportError and 0% test coverage.

Test plan

• Run optimization on get_files_for_language and verify generated tests no longer produce ImportError
• Verify existing optimization flows still work (the fallback to helper_functions ensures backward compat)

[claude]

PR Review Summary

Prek Checks

All checks passed. No formatting or linting issues found.

Mypy

No new type errors introduced by this PR. All mypy errors in the changed files are pre-existing.

Code Review

No critical issues found. The changes are well-structured:

• code_context_extractor.py: Correctly adds "statement" type to Jedi definition resolution, enabling module-level variables (e.g., _EXTENSION_REGISTRY) to be tracked as helpers. The all_helper_fqns construction properly combines first-degree and second-degree helper FQNs using a set for deduplication.
• models.py: Clean addition of testgen_helper_fqns: list[str] = [] with a safe default.
• function_optimizer.py: The fallback testgen_helper_fqns or [...] correctly preserves backward compatibility — if no FQNs are provided, it falls back to deriving them from helper_functions.
• Tests: Comprehensive updates to expected contexts reflecting the new statement-type helpers.

Test Coverage

File	Stmts	Miss	Coverage
codeflash/languages/python/context/code_context_extractor.py	624	54	91%
codeflash/models/models.py	627	138	78%
codeflash/optimization/function_optimizer.py	1164	949	18%
tests/test_code_context_extractor.py	1046	24	98%

• code_context_extractor.py (91%) — main file with the core logic changes — has strong coverage. The new "statement" type handling and all_helper_fqns construction are exercised by the updated tests.
• function_optimizer.py (18%) — low coverage is pre-existing; this file's changes are minimal (passing through testgen_helper_fqns), and it requires integration-level testing.
• No main branch comparison available (sandbox restriction), but the changes are additive and the test file at 98% coverage confirms thorough testing of the new behavior.

Codeflash Optimization PRs

No optimization PRs are mergeable — #1389 and #1291 both have failing CI checks (type-check-cli, pre-commit hooks).

---

Last updated: 2026-02-19T02:20Z