Refactor Python tests