Improve test cases for better coverage and reliability

[kernelshard]

This issue proposes improvements to the existing test cases in the repository to enhance code coverage, reliability, and maintainability.

Tasks:

• Review current test cases for gaps and inefficiencies
• Add missing test scenarios for edge cases and error handling
• Refactor and optimize existing tests for readability and performance
• Ensure all critical functionalities are covered
• Integrate updated tests into the CI pipeline

Expected Outcomes:

• Higher test coverage metrics
• Reduced flakiness in test runs
• Easier maintenance of test code

Feel free to suggest additional improvements or highlight neglected areas.