Updated unit tests to provide better coverage