Fix context check
Merge pull request #337 from my-personal-forks/fix-context-check

Fix context check