[PATCH 4/5] Re-write to comply with py.test format