Cleanup for pyflakes + name change in tests