cleaned up code to pass tests

cleaned up code to pass tests