NO-SNOW: Refactor test_np_unique into a class (#3943)

Author: sfc-gh-helmeleegy

tests/integ/modin/test_numpy.py

@@ -57,37 +57,44 @@ def test_np_may_share_memory():
         assert not np.may_share_memory(snow_df_A, native_df_A)
 
 
-@sql_count_checker(query_count=2)
-def test_np_unique():
-    # tests np.unique usage as seen in
-    # scikit-learn/sklearn/metrics/_ranking.py::average_precision_score
-    # and other places in the scikit-learn library
-    y_true_np = np.array([0, 0, 1, 1])
-    y_true_snow = pd.Series([0, 0, 1, 1])
-    res_np = np.unique(y_true_np)
-    res_snow = np.unique(y_true_snow)
-    assert (res_np == res_snow).all()
-
-    y_true_np_2d = np.array([[1, 2, 5, 6], [1, 2, 3, 4]])
-    y_true_snow_2d = pd.DataFrame({"a": [1, 2, 5, 6], "b": [1, 2, 3, 4]})
-    res_np = np.unique(y_true_np_2d)
-    res_snow = np.unique(y_true_snow_2d)
-    assert (res_np == res_snow).all()
-
-    # Verify that numpy throws type errors when we return NotImplemented
-    # when using optional parameters
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, return_index=True)
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, return_inverse=True)
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, return_counts=True)
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, axis=1)
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, equal_nan=False)
-    with pytest.raises(TypeError):
-        np.unique(y_true_snow_2d, sorted=False)
+class TestUnique:
+    @sql_count_checker(query_count=1)
+    def test_np_unique(self):
+        # tests np.unique usage as seen in
+        # scikit-learn/sklearn/metrics/_ranking.py::average_precision_score
+        # and other places in the scikit-learn library
+        numpy_ar = np.array([0, 0, 1, 1])
+        snow_ar = pd.Series([0, 0, 1, 1])
+        numpy_res = np.unique(numpy_ar)
+        snow_res = np.unique(snow_ar)
+        assert (numpy_res == snow_res).all()
+
+    @sql_count_checker(query_count=1)
+    def test_np_unique_2d(self):
+        numpy_ar_2d = np.array([[1, 2, 5, 6], [1, 2, 3, 4]])
+        snow_ar_2d = pd.DataFrame({"a": [1, 2, 5, 6], "b": [1, 2, 3, 4]})
+        numpy_res = np.unique(numpy_ar_2d)
+        snow_res = np.unique(snow_ar_2d)
+        assert (numpy_res == snow_res).all()
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            {"return_index": True},
+            {"return_inverse": True},
+            {"return_counts": True},
+            {"axis": 1},
+            {"equal_nan": False},
+            {"sorted": False},
+        ],
+    )
+    @sql_count_checker(query_count=0)
+    def test_np_unique_neg(self, kwargs):
+        snow_ar_2d = pd.DataFrame({"a": [1, 2, 5, 6], "b": [1, 2, 3, 4]})
+        # Verify that numpy throws type errors when we return NotImplemented
+        # when using optional parameters
+        with pytest.raises(TypeError):
+            np.unique(snow_ar_2d, **kwargs)
 
 
 def test_full_like():