JUnit tests updated