Adding unit test, cleanup etc.