Improve test coverage

[mikeqfu]

• Identify and add tests for untested or under-tested areas.
• Focus on critical or complex code paths.
• Ensure all new tests pass locally.
• Verify updated coverage on Codecov.