Fix py.test script