Added test case for unit of work