def test_get_records() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
records = data_table.get_records()
assert next(records) == {"a": 1, "b": 3}
assert next(records) == {"a": 2, "b": 4}
with pytest.raises(StopIteration):
next(records)
def test_custom_record(self):
data_table = DataTable(record_class=UserRecord)
data_table.add_record({"name": "Jon"})
data_table.add_record(UserRecord(name="test", age=12))
assert isinstance(data_table.get_record(0), UserRecord)
assert list(data_table.get_records()) == [
UserRecord(name="Jon"),
UserRecord(name="test", age=12),
]
with pytest.raises(ValueError):
data_table.add_record({"unknown": "Jon"})
with pytest.raises(ValueError):
data_table.add_record({})