From b9769f6b2d580c8187f437e7856bc9bb67fa40d6 Mon Sep 17 00:00:00 2001
From: "xingjun.wang" <xingjun.wxj@alibaba-inc.com>
Date: Thu, 22 Aug 2024 12:02:00 +0800
Subject: [PATCH] add ut for stream

---
 tests/msdatasets/test_stream_load.py | 49 ++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)
 create mode 100644 tests/msdatasets/test_stream_load.py

diff --git a/tests/msdatasets/test_stream_load.py b/tests/msdatasets/test_stream_load.py
new file mode 100644
index 00000000..3252d8cf
--- /dev/null
+++ b/tests/msdatasets/test_stream_load.py
@@ -0,0 +1,49 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import unittest
+
+from modelscope import MsDataset
+from modelscope.utils.logger import get_logger
+from modelscope.utils.test_utils import test_level
+
+logger = get_logger()
+
+
+class TestStreamLoad(unittest.TestCase):
+
+    def setUp(self):
+        ...
+
+    def tearDown(self):
+        ...
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_stream_read_zstd(self):
+        repo_id: str = 'swift/chinese-c4'
+        ds = MsDataset.load(repo_id, split='train', use_streaming=True)
+        sample = next(iter(ds))
+        logger.info(sample)
+
+        assert sample['url'], f'Failed to load sample from {repo_id}'
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_stream_imagefolder(self):
+        repo_id: str = 'wangxingjun778/test_new_dataset'
+        ds = MsDataset.load(repo_id, split='train', use_streaming=True)
+        sample = next(iter(ds))
+        logger.info(sample)
+
+        assert sample['image'], f'Failed to load sample from {repo_id}'
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_stream_parquet(self):
+        repo_id: str = 'swift/A-OKVQA'
+        ds = MsDataset.load(repo_id, split='train', use_streaming=True)
+        sample = next(iter(ds))
+        logger.info(sample)
+
+        assert sample['question'], f'Failed to load sample from {repo_id}'
+
+
+if __name__ == '__main__':
+    unittest.main()