Refactored test to a cleaner structure
Changed tests to py.test