wehub-resource-sync
41 lines
1008 B
Python
41 lines
1008 B
Python
import os
|
|
|
|
import fastavro
|
|
import pytest
|
|
|
|
import ray
|
|
|
|
schema = {
|
|
"type": "record",
|
|
"name": "TestRecord",
|
|
"fields": [{"name": "test_field", "type": "string"}],
|
|
}
|
|
|
|
|
|
def test_read_basic_avro_file(ray_start_regular_shared, tmp_path):
|
|
path = os.path.join(tmp_path, "sample.avro")
|
|
records = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
|
|
with open(path, "wb") as out:
|
|
fastavro.writer(out, schema, records)
|
|
|
|
ds = ray.data.read_avro(path)
|
|
|
|
expected = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
|
|
assert ds.take_all() == expected
|
|
|
|
|
|
def test_read_empty_avro_files(ray_start_regular_shared, tmp_path):
|
|
path = os.path.join(tmp_path, "empty.avro")
|
|
# Write an empty Avro file with the schema
|
|
with open(path, "wb") as out:
|
|
# Write the schema with no records
|
|
fastavro.writer(out, schema, [])
|
|
|
|
ds = ray.data.read_avro(path)
|
|
|
|
assert ds.count() == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main(["-v", __file__])
|