FIX-modin-project#2566: Check array shapes match for test_unique

Signed-off-by: Richard Lin <richard.lin.047@berkeley.edu>
richardlin047 · Dec 20, 2020 · 87df768 · 87df768
1 parent ba59173
commit 87df768
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/modin/pandas/test/test_general.py b/modin/pandas/test/test_general.py
@@ -528,10 +528,12 @@ def test_unique():
     modin_result = pd.unique([2, 1, 3, 3])
     pandas_result = pandas.unique([2, 1, 3, 3])
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(pd.Series([2] + [1] * 5))
     pandas_result = pandas.unique(pandas.Series([2] + [1] * 5))
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
         pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")])
@@ -540,6 +542,7 @@ def test_unique():
         pandas.Series([pandas.Timestamp("20160101"), pandas.Timestamp("20160101")])
     )
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
         pd.Series(
@@ -558,6 +561,7 @@ def test_unique():
         )
     )
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(
         pd.Index(
@@ -576,10 +580,12 @@ def test_unique():
         )
     )
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.unique(pd.Series(pd.Categorical(list("baabc"))))
     pandas_result = pandas.unique(pandas.Series(pandas.Categorical(list("baabc"))))
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
 
 @pytest.mark.parametrize("normalize, bins, dropna", [(True, 3, False)])

diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
@@ -3287,16 +3287,19 @@ def test_unique(data):
     modin_result = modin_series.unique()
     pandas_result = pandas_series.unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series([2, 1, 3, 3], name="A").unique()
     pandas_result = pandas.Series([2, 1, 3, 3], name="A").unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series([pd.Timestamp("2016-01-01") for _ in range(3)]).unique()
     pandas_result = pandas.Series(
         [pd.Timestamp("2016-01-01") for _ in range(3)]
     ).unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series(
         [pd.Timestamp("2016-01-01", tz="US/Eastern") for _ in range(3)]
@@ -3305,10 +3308,12 @@ def test_unique(data):
         [pd.Timestamp("2016-01-01", tz="US/Eastern") for _ in range(3)]
     ).unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pandas.Series(pd.Categorical(list("baabc"))).unique()
     pandas_result = pd.Series(pd.Categorical(list("baabc"))).unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
     modin_result = pd.Series(
         pd.Categorical(list("baabc"), categories=list("abc"), ordered=True)
@@ -3317,6 +3322,7 @@ def test_unique(data):
         pd.Categorical(list("baabc"), categories=list("abc"), ordered=True)
     ).unique()
     assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)