Add testcase
Move exception handling into context manager