Fix unit tests

A couple broken tests repaired