Update tests and make context optional again