Fixing Unit Testing Issues