Fix failing unit test, refactor python part