def test_get_reader_file_mode():
"""Test for getting a 'CsvReader' object with 'File' mode"""
source = "test/data/csv/mock_datasplitter_output"
with managed_env_var(
{os.path.basename(source): {
"TrainingInputMode": "File"
}}):
reader = _get_data(source=source)
assert isinstance(reader[0], mlio_file)
def test_get_reader_no_env_var():
"""Test for getting a 'CsvReader' object with no environmental variable"""
reader = _get_data(source="test/data/csv/mock_datasplitter_output")
assert isinstance(reader[0], mlio_file)
def test_get_reader_inmemory_mode():
"""Test for getting a 'CsvReader' object with 'InMemory' mode"""
buffer = BUFFER_DATA.encode()
reader = _get_data(source=buffer)
assert isinstance(reader[0], InMemoryStore)
def test_get_reader_mlio_file_object():
"""Test for getting a 'CsvReader' with a mlio.core.File object source"""
source = "test/data/csv/mock_datasplitter_output"
files = list_files(source, pattern="*")
reader = _get_data(source=files[0])
assert isinstance(reader[0], mlio_file)
def test_get_reader_pipe_mode():
"""Test for getting a 'CsvReader' object with 'Pipe' mode"""
with managed_env_var({"abc": {"TrainingInputMode": "Pipe"}}):
reader = _get_data(source="abc")
assert isinstance(reader[0], SageMakerPipe)
