| from tests.tools import RayUnittestBaseAsync, get_unittest_dataset_config |
| from trinity.buffer.buffer import get_buffer_reader |
| from trinity.buffer.reader import READER |
| from trinity.buffer.reader.file_reader import FileReader, TaskFileReader |
|
|
|
|
| @READER.register_module("custom") |
| class CustomReader(TaskFileReader): |
| """A custom reader for testing.""" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
|
|
|
|
| class TestBufferReader(RayUnittestBaseAsync): |
| async def test_buffer_reader_registration(self) -> None: |
| config = get_unittest_dataset_config("countdown", "train") |
| config.batch_size = 2 |
| config.storage_type = "custom" |
| reader = get_buffer_reader(config) |
| self.assertIsInstance(reader, CustomReader) |
| tasks = await reader.read_async() |
| self.assertEqual(len(tasks), 2) |
| config.storage_type = "file" |
| reader = get_buffer_reader(config) |
| self.assertIsInstance(reader, FileReader) |
| tasks = await reader.read_async() |
| self.assertEqual(len(tasks), 2) |
|
|
