Adding tests and fix for test fails & code cleanup