Spaces:
Runtime error
Runtime error
File size: 1,687 Bytes
e4f9cbe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
"""Tests for the JSON source."""
import json
import os
import pathlib
from ...schema import schema
from .json_source import ROW_ID_COLUMN, JSONDataset
from .source import SourceSchema
def test_simple_json(tmp_path: pathlib.Path) -> None:
json_records = [{'x': 1, 'y': 'ten'}, {'x': 2, 'y': 'twenty'}]
filename = 'test-dataset.jsonl'
filepath = os.path.join(tmp_path, filename)
with open(filepath, 'w') as f:
f.write(json.dumps(json_records))
source = JSONDataset(filepaths=[filepath])
source.setup()
source_schema = source.source_schema()
assert source_schema == SourceSchema(
fields=schema({
ROW_ID_COLUMN: 'int64',
'x': 'int64',
'y': 'string'
}).fields, num_items=2)
items = list(source.process())
assert items == [{
ROW_ID_COLUMN: 0,
'x': 1,
'y': 'ten'
}, {
ROW_ID_COLUMN: 1,
'x': 2,
'y': 'twenty'
}]
def test_simple_jsonl(tmp_path: pathlib.Path) -> None:
json_records = [{'x': 1, 'y': 'ten'}, {'x': 2, 'y': 'twenty'}]
json_lines = [json.dumps(record) + '\n' for record in json_records]
filename = 'test-dataset.jsonl'
filepath = os.path.join(tmp_path, filename)
with open(filepath, 'w') as f:
f.writelines(json_lines)
source = JSONDataset(dataset_name='test_dataset', filepaths=[filepath])
source.setup()
source_schema = source.source_schema()
assert source_schema == SourceSchema(
fields=schema({
ROW_ID_COLUMN: 'int64',
'x': 'int64',
'y': 'string'
}).fields, num_items=2)
items = list(source.process())
assert items == [{
ROW_ID_COLUMN: 0,
'x': 1,
'y': 'ten'
}, {
ROW_ID_COLUMN: 1,
'x': 2,
'y': 'twenty'
}]
|