refactor: Use provided context

[mitar]

Proposed changes

Context was not always passed through.

Checklist

• I have read the contributing guidelines and signed the CLA.
• I have read the security policy.
• I confirm that this pull request does not address a security
  vulnerability. If this pull request addresses a security vulnerability, I
  confirm that I got green light (please contact
  [email protected]) from the maintainers to push
  the changes.
• I have added tests that prove my fix is effective or that my feature works.
• I have added necessary documentation within the code base (if appropriate).