Refactor Flink status to modular architecture #4165
+1,183
−251
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nleigh
changed the title
Extracted data collection and formatting logic from the monolithic _print_flink_status_from_job_manager() function into reusable, testable helper functions in flink_tools.py. Changes: - Added collect_flink_job_details() to gather pod/job/resource info - Added format_flink_state_and_pods() to format state and counts - Added format_flink_jobs_table() to format jobs table - Added get_flink_job_name() helper for job name extraction - Refactored main function from ~290 lines to ~90 lines - Removed unused imports in status.py (shutil, groupby, get_runbook, FlinkJobs) - Maintained identical output and behavior This refactor: - Improves code maintainability and testability - Separates data collection from presentation logic - Makes functions reusable for future features - Reduces cognitive complexity of main function Related to FLINK-5725 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Added 10 test classes with 30+ test methods covering all newly extracted functions from the Flink status refactor: **Data Collection Tests:** - TestCollectFlinkJobDetails (3 tests) * Handles missing overview (stopped clusters) * Counts evicted pods correctly * Collects complete job/resource information **Formatting Tests:** - TestFormatFlinkStateAndPods (3 tests) - TestGetFlinkJobName (3 tests) - TestFormatFlinkJobsTable (5 tests) - TestGetFlinkInstanceDetails (2 tests) - TestFormatFlinkInstanceHeader (3 tests) - TestFormatFlinkInstanceMetadata (1 test) - TestFormatFlinkConfigLinks (1 test) - TestFormatFlinkLogCommands (1 test) - TestFormatFlinkMonitoringLinks (1 test) Test Coverage: - All new functions (collect_flink_job_details, format_flink_state_and_pods, get_flink_job_name, format_flink_jobs_table) - All existing formatting functions from first refactor (get_flink_instance_details, format_flink_instance_header, etc.) - Edge cases: stopped clusters, evicted pods, verbose modes, terminal sizing - All tests pass successfully via tox - All mocks use autospec=True per repo guidelines Benefits: - Prevents regressions in output formatting - Documents expected behavior - Makes future changes safer - Easy to test individual functions in isolation Related to FLINK-5725 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Improvements after self-review:
1. **Fix potential KeyError:**
- Changed pod["reason"] to pod.get("reason") to safely handle
Failed pods without a reason field
- Added test case for this edge case
2. **Add TypedDicts for better type safety:**
- PodCounts: Structured type for pod count statistics
- JobCounts: Structured type for job count statistics
- FlinkJobDetailsDict: Return type for collect_flink_job_details()
- FlinkInstanceDetails: Return type for get_flink_instance_details()
- Updated function signatures to use these types
Benefits:
- Better IDE autocomplete and type checking
- Prevents runtime errors from missing dictionary keys
- Makes function contracts more explicit
- Improves code documentation through types
All tests pass, all pre-commit checks pass.
Related to FLINK-5725
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <[email protected]>
nleigh
force-pushed
the
flink-status-refactor-modular
branch
from
6cd9a1d to
e38a0d1
Compare
Fixed type checking issues identified by mypy:
1. **TypedDict compatibility:**
- Added explicit type annotations to pod_counts and job_counts
- Changed pod_counts to: PodCounts = {...}
- Changed job_counts to: Optional[JobCounts] = None
2. **Verbose parameter type:**
- Changed format_flink_instance_header() parameter from bool to int
- Matches the actual usage where verbose can be 0, 1, 2, etc.
3. **Merge conflict resolution:**
- Updated format_flink_monitoring_links() to use CloudZero link
- Replaced Splunk link with CloudZero (from master PR #4152)
All checks pass:
- ✅ mypy: Success: no issues found
- ✅ All unit tests pass
- ✅ All pre-commit checks pass
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <[email protected]>
Updated test expectations to match the new modular output format:
1. **Fixed test helper function ordering:**
- _get_flink_base_status_verbose_1() now outputs in correct order:
Config SHA → Version → URL → Repos → Pool/Owner/Runbook
- Matches format_flink_instance_header + format_flink_instance_metadata
2. **Added URL to stopping state tests:**
- test_output_stopping_jobmanager
- test_output_stopping_taskmanagers
- URL is shown even for non-running states (from annotations)
3. **Fixed color code assertion:**
- test_format_stopped_state_with_evictions
- Check for "evicted" instead of "2 evicted" (handles ANSI colors)
4. **Fixed mock patch locations:**
- Changed get_team/get_runbook patches from flink_tools to monitoring_tools
- These functions are imported inside get_flink_instance_details
5. **Updated CloudZero link test:**
- Changed from splunk.yelpcorp.com to app.cloudzero.com
All tests now pass (11 status tests + 36 flink_tools tests)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <[email protected]>
nemacysts
reviewed
Nov 20, 2025
Member
nemacysts
left a comment
nemacysts
left a comment
There was a problem hiding this comment.
ty - i've been meaning to try to shrink status.py a bit and move the operator-specific logic out as much as possible :)
paasta_tools/flink_tools.py
Outdated
Comment on lines
684
to
689
| state: str, | ||
| pod_counts: PodCounts, | ||
| job_counts: Optional[JobCounts], | ||
| taskmanagers: Optional[int], | ||
| slots_available: Optional[int], | ||
| slots_total: Optional[int], |
Member
There was a problem hiding this comment.
hmm, job_details is already a TypedDict - would we lose anything if we just passed job_details as a whole rather than each attribute individually?
tests/test_flink_tools.py
Outdated
Comment on lines
344
to
353
| assert result["state"] == "stopped" | ||
| assert result["pod_counts"]["running"] == 2 | ||
| assert result["pod_counts"]["evicted"] == 0 | ||
| assert result["pod_counts"]["other"] == 0 | ||
| assert result["pod_counts"]["total"] == 2 | ||
| assert result["job_counts"] is None | ||
| assert result["taskmanagers"] is None | ||
| assert result["slots_available"] is None | ||
| assert result["slots_total"] is None | ||
| assert result["jobs"] == [] |
Member
There was a problem hiding this comment.
should we construct an expected_result dict and compare against that rather than manually comparing each key?
tests/test_flink_tools.py
Outdated
Comment on lines
366
to
373
| overview = mock.Mock(autospec=True) | ||
| overview.jobs_running = 1 | ||
| overview.jobs_finished = 2 | ||
| overview.jobs_failed = 0 | ||
| overview.jobs_cancelled = 1 | ||
| overview.taskmanagers = 3 | ||
| overview.slots_available = 5 | ||
| overview.slots_total = 15 |
Member
There was a problem hiding this comment.
overview is a typeddict, no? do we need a mock here? can't we just construct the expected dict?
tests/test_flink_tools.py
Outdated
| overview.slots_available = 10 | ||
| overview.slots_total = 25 | ||
|
|
||
| mock_jobs = [mock.Mock(autospec=True), mock.Mock(autospec=True)] |
Member
There was a problem hiding this comment.
I don't think we want autospec=True here - i think we want spec=$CLASS_OR_OBJECT or to use https://docs.python.org/3/library/unittest.mock.html#unittest.mock.create_autospec
Replace Any type annotations with concrete types: - flink_config: Optional[Any] -> Optional[FlinkConfig] - flink_instance_config: Any -> FlinkDeploymentConfig This addresses review comment from PR #4165. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Move get_runbook and get_team imports from inside get_flink_instance_details function to the top-level imports section, following Python best practices. Update test mocks to patch the functions where they are used (flink_tools module) rather than where they are defined (monitoring_tools module), which is required when imports are at module level. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Remove code that strips 'config' prefix from config_sha to make it clearer that this is not a git SHA. Now displays 'config123456' instead of '123456'. This addresses non-blocking review feedback from PR #4165. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Restore the validation check that ensures config_sha is present in the Flink metadata labels. This prevents silently ignoring missing config_sha which could indicate a serious configuration issue. Previously the validation was in status.py but was removed during the refactoring. Now it's properly placed in get_flink_instance_details() where it validates the input early. This addresses review feedback from PR #4165. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Consolidate two separate 'if status["state"] == "running"' blocks into a single block that fetches flink_config, overview, and jobs together. This improves code clarity by grouping all running-state API calls in one place instead of having them split across the function. Also update test expectations to reflect the 'config' prefix being kept in config_sha display (from previous commit). This addresses review feedback from PR #4165. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Address Feedback
Address Feedback class FlinkJobDetailsDict(TypedDict, total=False): "just curious: why is this non-total? seems like we only have one function that returns this"
Address Feedback 1. format_flink_instance_header - Changed parameter type from Mapping[str, Any] to FlinkInstanceDetails for better type safety. 2. format_flink_instance_metadata - Changed parameter type from Mapping[str, Any] to FlinkInstanceDetails for better type safety. 3. format_flink_state_and_pods - Simplified to take FlinkJobDetailsDict as a single parameter instead of 6 individual parameters. The function now extracts the values it needs internally. 4. Call site in status.py - Simplified from passing 6 parameters to just passing job_details. 5. Tests for TestFormatFlinkStateAndPods - Updated to construct FlinkJobDetailsDict dicts instead of passing individual parameters.
Address feedback
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
- Convert verbose int to bool for format_flink_instance_header call - Fix extra space in taskmanagers output line 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Restores the original comment explaining why we validate config_sha early and raise ValueError if missing. This addresses PR review feedback asking if the early exit behavior was preserved.
The function names format_flink_jobs_table and append_pod_status are self-explanatory; the comments don't add value.
The parenthetical list of fields would get out of date as metadata is added or changed.
Same reasoning as the previous commit - avoids staleness.
Compare against expected dict instead of asserting each key individually, making tests more readable and maintainable.
Test that get_flink_instance_details raises ValueError when config_sha label is missing from metadata, ensuring proper error handling for corrupted/invalid Flink cluster states.
Test that collect_flink_job_details handles status dicts without a pod_status key, returning zero counts instead of crashing during cluster transitions.
Split the single try/except block into separate blocks for each API call (config, overview, jobs). This provides more specific error messages so users can identify which API call failed.
Extract phase via pod.get("phase") instead of direct key access
to prevent KeyError if a malformed pod entry is missing the
phase field.
The function is now only defined in flink_tools.py. Updated test import to use the flink_tools version.
Return early with error message if flink metadata is missing, preventing AttributeError when passing None to get_flink_instance_details.
Contributor
Author
|
@nemacysts Hey, thanks for all the review feedback 🙌 Example Split |
Member
|
@nleigh splitting this up would be nice - reviewing this much generated code is slightly painful (which is why i keep putting it off :p) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Plan to add new features(#4162) but found the current flink paasta status code made it difficult to add/review new changes, refactored existing to be more modular.
Code Refactoring:
paasta_tools/flink_tools.pyNew Functions (Data Collection):
get_flink_instance_details()- Collects metadata, version, pool, team,runbook
collect_flink_job_details()- Collects job, pod, and resource informationNew Functions (Formatting):
format_flink_instance_header()- Formats config SHA, version, URLformat_flink_instance_metadata()- Formats repo links, pool, owner, runbookformat_flink_config_links()- Formats yelpsoa/srv-configs linksformat_flink_log_commands()- Formats paasta logs commandsformat_flink_monitoring_links()- Formats Grafana/cost linksformat_flink_state_and_pods()- Formats state, pods, jobs summaryformat_flink_jobs_table()- Formats jobs table with proper columnsget_flink_job_name()- Helper to extract job nameType Safety Improvements:
PodCounts- Pod count statisticsJobCounts- Job count statisticsFlinkJobDetailsDict- Collected job detailsFlinkInstanceDetails- Instance metadataSame output as existing (Tested 9th Jan after recent commits)