Adding support for pytest