Forgot to commit unit test updates