Added a test for ArangeDataset