Fix context handling regression (yay tests)