Adjusted code to make unit tests work.