improving python unit tests