Fixed dataset method and added a couple of unit tests