Fixing unit test.  Still needs more work but unbreaks things.