chore: follow up context.Context support

[johejo]

Follow up context support to clean up codes,
ref: https://blog.golang.org/module-compatibility#TOC_2.

Checklist

• I acknowledge that all my contributions will be made under the project's license
• I have made a material change to the repo (functionality, testing, spelling, grammar)
• I have read the Contribution Guidelines and my PR follows them
• I have titled the PR appropriately
• I have updated my branch with the main branch
• I have added tests that prove my fix is effective or that my feature works
• I have added the necessary documentation about the functionality in the appropriate .md file
• I have added inline documentation to the code I modified

If you have questions, please file a support ticket, or create a GitHub Issue in this repository.

[JenniferMah]

Thanks @johejo!