Skip to content

Commit 0c70d7c

Browse files
committed
add tests
1 parent 81be345 commit 0c70d7c

File tree

2 files changed

+149
-0
lines changed

2 files changed

+149
-0
lines changed

protocol/tests/conftest.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import pytest
2+
import pandas as pd
3+
4+
5+
@pytest.fixture
6+
def constructor_frame(data):
7+
return pd.DataFrame(data)

protocol/tests/test_protocol.py

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
import pytest
2+
import numpy as np
3+
from conftest import *
4+
5+
6+
@pytest.mark.parametrize("test_data",
7+
[
8+
({'a': [np.array([1, 2, 3]), np.array([4, 5, 6])],
9+
'b': [np.array([1.5, 2.0, 3.2]), np.array([4.1, 5.7, 6.9])]},
10+
np.object_, None),
11+
({'a': [1.5, 2.5, 3.5], 'b': [9.2, 10.5, 11.8]}, np.float64, None),
12+
({'A': [1, 2, 3, 4], 'B': [1, 2, 3, 4]}, np.int64, np.float64)
13+
],
14+
ids=["array_data", "float_data", "int_data"])
15+
def test_only_one_data(test_data):
16+
data, dtype, new_dtype = test_data
17+
columns = list(data.keys())
18+
df = constructor_frame(data)
19+
df2 = df.__dataframe__()
20+
new_dtype = dtype if new_dtype is None else new_dtype
21+
assert df.columns.values.tolist() == columns
22+
val = len(df[columns[0]])-1
23+
column_size = df.size
24+
for column in columns:
25+
assert df[column].tolist() == df[column].tolist()
26+
assert df[column].dtype.type is dtype
27+
assert df2.get_column_by_name(column).null_count == 0
28+
assert df2.get_column_by_name(column).size == column_size
29+
assert df2.get_column_by_name(column).offset == 0
30+
assert not df2["x"].is_masked
31+
n = np.random.randint(0, val)
32+
(df[column])[n] = None
33+
assert df[column].dtype.type is new_dtype
34+
assert df2.get_column_by_name(column).null_count == 1
35+
36+
37+
def test_float_int():
38+
df = constructor_frame({'a': [1, 2, 3], 'b': [3, 4, 5],
39+
'c': [1.5, 2.5, 3.5], 'd': [9, 10, 11]})
40+
df2 = df.__dataframe__()
41+
columns = ['a', 'b', 'c', 'd']
42+
assert df.columns.values.tolist() == columns
43+
for column in columns:
44+
assert df[column].tolist() == df[column].tolist()
45+
if column is 'c':
46+
assert df[column].dtype.type is np.float64
47+
else:
48+
assert df[column].dtype.type is np.int64
49+
50+
assert df2.get_column_by_name(column).null_count == 0
51+
assert df2.get_column_by_name(column).size == 3
52+
assert df2.get_column_by_name(column).offset == 0
53+
54+
n = np.random.randint(0, 2)
55+
(df[column])[n] = None
56+
assert df[column].dtype.type is np.float64
57+
assert df2.get_column_by_name(column).null_count == 1
58+
59+
60+
def test_mixed_intfloatbool():
61+
df = constructor_frame({"x": np.array([True, True, False]),
62+
"y": np.array([1, 2, 0]),
63+
"z": np.array([9.2, 10.5, 11.8])})
64+
df2 = df.__dataframe__()
65+
columns = ['x', 'y', 'z']
66+
assert df.columns.values.tolist() == columns
67+
for column in columns:
68+
assert df[column].tolist() == df[column].tolist()
69+
assert df2.get_column_by_name(column).null_count == 0
70+
assert df2.get_column_by_name(column).size == 3
71+
assert df2.get_column_by_name(column).offset == 0
72+
73+
assert df["x"].dtype.type is np.bool_
74+
assert df["y"].dtype.type is np.int32
75+
assert df["z"].dtype.type is np.float64
76+
77+
assert df2.get_column_by_name("x")._allow_copy == True
78+
79+
for column in columns:
80+
n = np.random.randint(0, 2)
81+
(df[column])[n] = None
82+
if column is "x":
83+
assert df[column].dtype.type is np.object_
84+
else:
85+
assert df[column].dtype.type is np.float64
86+
assert df2.get_column_by_name(column).null_count == 1
87+
88+
89+
def test_string_dtype():
90+
df = constructor_frame({"A": ["a", "b", "cdef", "", "g"]})
91+
df2 = df.__dataframe__()
92+
columns = ['A']
93+
assert df.columns.values.tolist() == columns
94+
for column in columns:
95+
assert df[column].tolist() == df[column].tolist()
96+
assert df[column].dtype.type is np.object_
97+
assert df2.get_column_by_name(column).null_count == 0
98+
99+
100+
def test_categorical():
101+
df = constructor_frame({"year": [2012, 2013, 2015, 2019], "weekday": [0, 1, 4, 6]})
102+
df = df.categorize("year", min_value=2012, max_value=2019)
103+
df = df.categorize("weekday", labels=["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"])
104+
# Some detailed testing for correctness of dtype and null handling:
105+
col = df.__dataframe__().get_column_by_name("year")
106+
assert col.describe_categorical == (False, True, {0: 2012, 1: 2013, 2: 2014, 3: 2015, 4: 2016, 5: 2017, 6: 2018, 7: 2019})
107+
assert col.describe_null == (0, None)
108+
col2 = df.__dataframe__().get_column_by_name("weekday")
109+
assert col2.describe_categorical == (False, True, {0: "Mon", 1: "Tue", 2: "Wed", 3: "Thu", 4: "Fri", 5: "Sat", 6: "Sun"})
110+
assert col2.describe_null == (0, None)
111+
112+
113+
def test_dataframe():
114+
df = constructor_frame({"x": [True, True, False], "y": [1, 2, 0], "z": [9.2, 10.5, 11.8]})
115+
df2 = df.__dataframe__()
116+
assert df2._allow_copy == True
117+
assert df2.num_columns() == 3
118+
assert df2.num_rows() == 3
119+
assert df2.num_chunks() == 1
120+
assert df2.column_names() == ["x", "y", "z"]
121+
assert df2.select_columns((0, 2))._df[:, 0].tolist() == df2.select_columns_by_name(("x", "z"))._df[:, 0].tolist()
122+
assert df2.select_columns((0, 2))._df[:, 1].tolist() == df2.select_columns_by_name(("x", "z"))._df[:, 1].tolist()
123+
124+
125+
def test_chunks():
126+
df = constructor_frame({"x": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
127+
df2 = df.__dataframe__()
128+
chunk_iter = iter(df2.get_chunks(3))
129+
chunk = next(chunk_iter)
130+
assert chunk.num_rows() == 4
131+
chunk = next(chunk_iter)
132+
assert chunk.num_rows() == 4
133+
chunk = next(chunk_iter)
134+
assert chunk.num_rows() == 2
135+
with pytest.raises(StopIteration):
136+
chunk = next(chunk_iter)
137+
138+
139+
def test_get_chunks():
140+
df = constructor_frame({"x": [1]})
141+
df2 = df.__dataframe__()
142+
assert df2.get_chunks() == 1

0 commit comments

Comments
 (0)