factor out test code and make the python wrapper more generic