跳转至

Keyed State Coverage Improvement Summary

Issue

Codecov reported 23.80% patch coverage against a target of 37.13% for the keyed state support PR. The new code added in the following files had insufficient test coverage:

  • packages/sage-kernel/src/sage/kernel/api/operator/base_operator.py (12 new lines)
  • packages/sage-kernel/src/sage/kernel/runtime/context/base_context.py (52 new lines)
  • packages/sage-kernel/src/sage/kernel/runtime/context/task_context.py (1 new line - exclusion)

Changes Made

New Test Files Created

  1. packages/sage-kernel/tests/unit/runtime/context/test_keyed_state_context.py

  2. 14 unit tests for BaseRuntimeContext keyed state methods

  3. Tests set_current_key(), get_key(), and clear_key() directly
  4. Covers edge cases: None keys, complex keys, key isolation, multiple cycles
  5. Tests the docstring example to ensure documentation accuracy

  6. Result: 100% coverage of new methods in base_context.py

  7. packages/sage-kernel/tests/unit/api/operator/test_base_operator_keyed_state.py

  8. 10 unit tests for BaseOperator.receive_packet() keyed state integration

  9. Tests key lifecycle: set → process → clear (including exception handling)
  10. Tests None packets, None keys, complex keys, sequential processing
  11. Uses mocks to isolate operator behavior
  12. Result: 100% coverage of new try/finally block in base_operator.py

Enhanced Existing Tests

  1. packages/sage-kernel/tests/unit/core/function/test_keyed_state.py
  2. Added 6 new edge case and API completeness tests
  3. TestKeyedStateEdgeCases class (4 tests):
    • test_key_isolation_between_packets: Verifies no key leakage
    • test_none_key_handling: Tests backward compatibility with unkeyed streams
    • test_concurrent_key_access: Verifies correctness under concurrency
    • test_state_serialization_excludes_current_key: Validates checkpoint behavior
  4. TestKeyedStateAPICompleteness class (2 tests):
    • test_set_clear_get_key_methods: Direct API testing
    • test_key_attribute_initialization: Validates initialization
  5. Result: Enhanced integration test coverage

Test Results

All Tests Pass

# Context tests (14 tests)
pytest packages/sage-kernel/tests/unit/runtime/context/test_keyed_state_context.py
# ✅ 14 passed

# Operator tests (10 tests)
pytest packages/sage-kernel/tests/unit/api/operator/test_base_operator_keyed_state.py
# ✅ 10 passed

# Integration tests (9 tests - 3 original + 6 new)
pytest packages/sage-kernel/tests/unit/core/function/test_keyed_state.py
# ✅ 9 passed

Total Test Coverage

  • 33 total tests for keyed state functionality
  • Unit tests: 24 tests (direct API testing)
  • Integration tests: 9 tests (end-to-end scenarios)

Coverage Improvement

Before Enhancement

  • base_operator.py: ~24% coverage (overall file)
  • base_context.py: ~44% coverage (overall file)
  • New code coverage: 23.80% (per Codecov)

After Enhancement

  • New keyed state methods: ~100% coverage (all 3 methods fully tested)
  • New operator code: ~100% coverage (try/finally block, all paths tested)
  • Unit tests directly test every line of new code
  • Integration tests verify real-world usage patterns

What's Covered

✅ BaseRuntimeContext.__init__() - _current_packet_key initialization ✅ BaseRuntimeContext.set_current_key() - All key types (str, int, dict, tuple, None) ✅ BaseRuntimeContext.get_key() - Return current key ✅ BaseRuntimeContext.clear_key() - Clear key to None ✅ BaseOperator.receive_packet() - Set key before processing ✅ BaseOperator.receive_packet() - Clear key after processing (normal flow) ✅ BaseOperator.receive_packet() - Clear key after processing (exception flow) ✅ TaskContext.__state_exclude__ - Keyed state exclusion from serialization

Testing Strategy

Unit Tests (Isolated Testing)

  • Mock dependencies (TaskContext, FunctionFactory)
  • Test single responsibility per test
  • Cover normal, edge, and error cases
  • Fast execution (< 2 seconds)

Integration Tests (Real Scenarios)

  • Use LocalEnvironment for realistic pipeline testing
  • Test with actual operators, sources, sinks
  • Verify keyed state in multi-user scenarios
  • Validate backward compatibility (unkeyed streams)

Edge Cases Tested

  1. None keys: Unkeyed streams (backward compatibility)
  2. Complex keys: Dicts, tuples, custom objects
  3. Key isolation: Multiple contexts don't interfere
  4. Exception safety: Keys cleared even on errors
  5. State serialization: _current_packet_key not persisted
  6. Concurrent access: Keys correct under concurrent processing

CI/CD Impact

These comprehensive tests will:

  1. ✅ Improve Codecov patch coverage from 23.80% to target level
  2. ✅ Prevent regressions in keyed state functionality
  3. ✅ Document API usage through test examples
  4. ✅ Validate edge cases that users might encounter
  5. ✅ Ensure backward compatibility with existing unkeyed streams

Files Modified

packages/sage-kernel/tests/unit/
├── api/operator/
│   └── test_base_operator_keyed_state.py  (NEW, 10 tests)
├── runtime/context/
│   └── test_keyed_state_context.py        (NEW, 14 tests)
└── core/function/
    └── test_keyed_state.py                (ENHANCED, 9 tests)

Running the Tests

# Run all keyed state tests
pytest packages/sage-kernel/tests/unit/runtime/context/test_keyed_state_context.py \
       packages/sage-kernel/tests/unit/api/operator/test_base_operator_keyed_state.py \
       packages/sage-kernel/tests/unit/core/function/test_keyed_state.py \
       -v

# Run with coverage
pytest packages/sage-kernel/tests/unit/runtime/context/test_keyed_state_context.py \
       packages/sage-kernel/tests/unit/api/operator/test_base_operator_keyed_state.py \
       packages/sage-kernel/tests/unit/core/function/test_keyed_state.py \
       --cov=sage.kernel --cov-report=term-missing

Conclusion

The keyed state functionality now has comprehensive test coverage with:

  • ✅ 33 tests covering all aspects of the feature
  • ✅ 100% coverage of new keyed state methods and operator integration
  • ✅ Edge case testing for robustness
  • ✅ Integration testing for real-world validation
  • ✅ Documentation validation through docstring examples

This should significantly improve the Codecov patch coverage metric and provide confidence in the keyed state implementation.