Small bug fixes on unit tests