Skip to content

Commit 0a2ddd6

Browse files
committed
fix unittest
1 parent c526a17 commit 0a2ddd6

File tree

1 file changed

+10
-6
lines changed

1 file changed

+10
-6
lines changed

tests/tools/test_process_data.py

+10-6
Original file line numberDiff line numberDiff line change
@@ -182,12 +182,16 @@ def test_ray_image(self):
182182

183183
self.assertTrue(osp.exists(tmp_out_path))
184184

185-
import ray
186-
res_ds = ray.data.read_json(tmp_out_path)
187-
res_ds = res_ds.to_pandas().to_dict(orient='records')
188-
189-
self.assertEqual(len(res_ds), 3)
190-
for item in res_ds:
185+
from datasets import load_dataset
186+
jsonl_files = [os.path.join(tmp_out_path, f) \
187+
for f in os.listdir(tmp_out_path) \
188+
if f.endswith('.json')]
189+
dataset = load_dataset(
190+
'json',
191+
data_files={'jsonl': jsonl_files})
192+
193+
self.assertEqual(len(dataset['jsonl']), 3)
194+
for item in dataset['jsonl']:
191195
self.assertIn('aspect_ratios', item['__dj__stats__'])
192196

193197

0 commit comments

Comments
 (0)