fix and update unit tests