26 lines
650 B
Python
26 lines
650 B
Python
|
import pandas as pd
|
||
|
import pytest
|
||
|
from src.cleaner.dataframe import build_df
|
||
|
|
||
|
|
||
|
@pytest.fixture
|
||
|
def sample_objs():
|
||
|
return [
|
||
|
{'id': 1, 'name': 'Alice', 'age': 25},
|
||
|
{'id': 2, 'name': 'Bob', 'age': 30},
|
||
|
{'id': 3, 'name': 'Charlie', 'age': 22}
|
||
|
]
|
||
|
|
||
|
|
||
|
def test_build_df(sample_objs):
|
||
|
result_df = build_df(sample_objs)
|
||
|
|
||
|
assert isinstance(result_df, pd.DataFrame)
|
||
|
assert set(result_df.columns) == {'name', 'age'}
|
||
|
|
||
|
|
||
|
def test_build_df_missing_id_column():
|
||
|
objs_missing_id = [{'name': 'Alice', 'age': 25}, {'name': 'Bob', 'age': 30}]
|
||
|
with pytest.raises(KeyError, match="'id'"):
|
||
|
build_df(objs_missing_id)
|