Adjusted code to allow unit tests to pass