bug fix, and cleaning up unit tests