#100 Use pandas testing functions for tests

Author: astropenguin

tests/test_asdata.py

@@ -4,8 +4,9 @@
 
 # dependencies
 import pandas as pd
+from pandas.testing import assert_frame_equal, assert_series_equal
 from data import Weather, weather, df_weather_true, ser_weather_true
-from pandas_dataclasses import asdataframe, asseries, Spec
+from pandas_dataclasses import Spec, asdataframe, asseries
 from pandas_dataclasses.core.asdata import get_attrs, get_columns, get_data, get_index
 
 
@@ -15,92 +16,45 @@
 
 # test functions
 def test_asseries() -> None:
-    ser_weather = asseries(weather)
-
-    assert ser_weather.attrs == ser_weather_true.attrs
-    assert ser_weather.name == ser_weather_true.name
-    assert ser_weather.dtype == ser_weather_true.dtype
-    assert (ser_weather == ser_weather_true).all()
-
-    assert ser_weather.index.name == ser_weather_true.index.name
-    assert ser_weather.index.dtype == ser_weather_true.index.dtype
-    assert (ser_weather.index == ser_weather_true.index).all()
+    assert_series_equal(asseries(weather), ser_weather_true)
 
 
 def test_asdataframe() -> None:
-    df_weather = asdataframe(weather)
-
-    assert df_weather.attrs == df_weather_true.attrs
-    assert df_weather.iloc[:, 0].dtype == df_weather_true.iloc[:, 0].dtype
-    assert df_weather.iloc[:, 1].dtype == df_weather_true.iloc[:, 1].dtype
-    assert df_weather.iloc[:, 2].dtype == df_weather_true.iloc[:, 2].dtype
-    assert df_weather.iloc[:, 3].dtype == df_weather_true.iloc[:, 3].dtype
-    assert (df_weather == df_weather_true).all().all()
-
-    assert df_weather.columns.names == df_weather_true.columns.names
-    assert (df_weather.columns.dtypes == df_weather_true.columns.dtypes).all()
-    assert (df_weather.columns == df_weather_true.columns).all()
-
-    assert df_weather.index.names == df_weather_true.index.names
-    assert (df_weather.index.dtypes == df_weather_true.index.dtypes).all()  # type: ignore
-    assert (df_weather.index == df_weather_true.index).all()
+    assert_frame_equal(asdataframe(weather), df_weather_true)
 
 
 def test_get_attrs() -> None:
     attrs = get_attrs(spec)
-    keys = list(attrs.keys())
-    values = list(attrs.values())
 
-    assert keys[0] == spec.fields.of_attr[0].name
-    assert keys[1] == spec.fields.of_attr[1].name
-    assert keys[2] == spec.fields.of_attr[2].name
-
-    assert values[0] == spec.fields.of_attr[0].default
-    assert values[1] == spec.fields.of_attr[1].default
-    assert values[2] == spec.fields.of_attr[2].default
+    for i, (key, val) in enumerate(attrs.items()):
+        assert key == spec.fields.of_attr[i].name
+        assert val == spec.fields.of_attr[i].default
 
 
 def test_get_columns() -> None:
     columns = cast(pd.Index, get_columns(spec))
 
-    assert columns.names[0] == spec.fields.of_column[0].name
-    assert columns.names[1] == spec.fields.of_column[1].name
+    for i in range(len(columns)):
+        assert columns[i] == spec.fields.of_data[i].name
 
-    assert columns[0] == spec.fields.of_data[0].name
-    assert columns[1] == spec.fields.of_data[1].name
-    assert columns[2] == spec.fields.of_data[2].name
-    assert columns[3] == spec.fields.of_data[3].name
+    for i in range(columns.nlevels):
+        assert columns.names[i] == spec.fields.of_column[i].name
 
 
 def test_get_data() -> None:
     data = get_data(spec)
-    keys = list(data.keys())
-    values = list(data.values())
-
-    assert keys[0] == spec.fields.of_data[0].name
-    assert keys[1] == spec.fields.of_data[1].name
-    assert keys[2] == spec.fields.of_data[2].name
-    assert keys[3] == spec.fields.of_data[3].name
-
-    assert values[0].dtype.name == spec.fields.of_data[0].dtype
-    assert values[1].dtype.name == spec.fields.of_data[1].dtype
-    assert values[2].dtype.name == spec.fields.of_data[2].dtype
-    assert values[3].dtype.name == spec.fields.of_data[3].dtype
 
-    assert (values[0] == spec.fields.of_data[0].default).all()
-    assert (values[1] == spec.fields.of_data[1].default).all()
-    assert (values[2] == spec.fields.of_data[2].default).all()
-    assert (values[3] == spec.fields.of_data[3].default).all()
+    for i, (key, val) in enumerate(data.items()):
+        assert key == spec.fields.of_data[i].name
+        assert val.dtype.name == spec.fields.of_data[i].dtype
+        assert (val == spec.fields.of_data[i].default).all()
 
 
 def test_get_index() -> None:
     index = cast(pd.Index, get_index(spec))
 
-    assert index.names[0] == spec.fields.of_index[0].name
-    assert index.names[1] == spec.fields.of_index[1].name
-
-    assert index.get_level_values(0).dtype == spec.fields.of_index[0].dtype
-    assert index.get_level_values(1).dtype == spec.fields.of_index[1].dtype
-
-    assert (index.get_level_values(0) == spec.fields.of_index[0].default).all()
-    assert (index.get_level_values(1) == spec.fields.of_index[1].default).all()
+    for i in range(index.nlevels):
+        level = index.get_level_values(i)
+        assert level.name == spec.fields.of_index[i].name
+        assert level.dtype.name == spec.fields.of_index[i].dtype
+        assert (level == spec.fields.of_index[i].default).all()

tests/test_mixins.py

@@ -4,6 +4,7 @@
 
 # dependencies
 import pandas as pd
+from pandas.testing import assert_frame_equal, assert_series_equal
 from data import Weather, weather, df_weather_true, ser_weather_true
 from pandas_dataclasses import As, AsDataFrame, AsSeries
 
@@ -13,7 +14,7 @@ class CustomDataFrame(pd.DataFrame):
     pass
 
 
-class CustomSeries(pd.Series):
+class CustomSeries(pd.Series):  # type: ignore
     pass
 
 
@@ -49,7 +50,7 @@ def test_dataframe_weather() -> None:
     )
 
     assert isinstance(df_weather, pd.DataFrame)
-    assert (df_weather == df_weather_true).all().all()
+    assert_frame_equal(df_weather, df_weather_true)
 
 
 def test_custom_dataframe_weather() -> None:
@@ -63,7 +64,7 @@ def test_custom_dataframe_weather() -> None:
     )
 
     assert isinstance(df_weather, CustomDataFrame)
-    assert (df_weather == df_weather_true).all().all()
+    assert_frame_equal(df_weather, df_weather_true, check_frame_type=False)
 
 
 def test_series_weather() -> None:
@@ -77,7 +78,7 @@ def test_series_weather() -> None:
     )
 
     assert isinstance(ser_weather, pd.Series)
-    assert (ser_weather == ser_weather_true).all()
+    assert_series_equal(ser_weather, ser_weather_true)
 
 
 def test_custom_series_weather() -> None:
@@ -91,4 +92,4 @@ def test_custom_series_weather() -> None:
     )
 
     assert isinstance(ser_weather, CustomSeries)
-    assert (ser_weather == ser_weather_true).all()
+    assert_series_equal(ser_weather, ser_weather_true, check_series_type=False)