JUnit added + minor changes