Use context manager to protect test execution