Added junit and some changes to build