Fix test, add new test, clean stuff up