Added required code for tests to run