Add test for nsorted method

Author: MartinBraquet

pandas/tests/frame/methods/test_nlargest.py

@@ -75,8 +75,8 @@ def test_nsorted_n(self, nselect_method, n: int, columns):
         )
         if "b" in columns:
             error_msg = (
-                f"Column 'b' has dtype (object|str), "
-                f"cannot use method '{nselect_method}' with this dtype"
+                "Column 'b' has dtype (object|str), "
+                "cannot use n-sorting with this dtype"
             )
             with pytest.raises(TypeError, match=error_msg):
                 getattr(df, nselect_method)(n, columns)
@@ -87,6 +87,29 @@ def test_nsorted_n(self, nselect_method, n: int, columns):
             expected = df.sort_values(columns, ascending=ascending).head(n)
             tm.assert_frame_equal(result, expected)
 
+    def test_nsorted(self):
+        df = pd.DataFrame(
+            {
+                "x": [2, 2, 1],
+                "y": [3, 2, 1],
+            },
+            index=["a", "b", "c"],
+        )
+        cols = ["x", "y"]
+        ascending = [True, False]
+        n = 2
+        df_sort_values = df.sort_values(cols, ascending=ascending).head(n)
+        result = df.nsorted(n, cols, ascending=ascending)
+        tm.assert_frame_equal(result, df_sort_values)
+        expected = pd.DataFrame(
+            {
+                "x": [1, 2],
+                "y": [1, 3],
+            },
+            index=["c", "a"],
+        )
+        tm.assert_frame_equal(result, expected)
+
     @pytest.mark.parametrize(
         "columns", [["group", "category_string"], ["group", "string"]]
     )
@@ -95,7 +118,7 @@ def test_nsorted_error(self, df_main_dtypes, nselect_method, columns):
         col = columns[1]
         error_msg = (
             f"Column '{col}' has dtype {df[col].dtype}, "
-            f"cannot use method '{nselect_method}' with this dtype"
+            f"cannot use n-sorting with this dtype"
         )
         # escape some characters that may be in the repr
         error_msg = (

pandas/tests/series/methods/test_nlargest.py

@@ -20,7 +20,7 @@ def assert_check_nselect_boundary(vals, dtype, method):
     tm.assert_series_equal(result, expected)
 
 
-class TestSeriesNLargestNSmallest:
+class TestSeriesNSorted:
     @pytest.mark.parametrize(
         "r",
         [
@@ -37,7 +37,7 @@ class TestSeriesNLargestNSmallest:
     @pytest.mark.parametrize("arg", [2, 5, 0, -1])
     def test_nlargest_error(self, r, method, arg):
         dt = r.dtype
-        msg = f"Cannot use method 'n(largest|smallest)' with dtype {dt}"
+        msg = f"Cannot use n-sorting with dtype {dt}"
         with pytest.raises(TypeError, match=msg):
             getattr(r, method)(arg)
 
@@ -78,6 +78,9 @@ def test_nsmallest_nlargest(self, data):
         tm.assert_series_equal(ser.nlargest(len(ser)), ser.iloc[[4, 0, 1, 3, 2]])
         tm.assert_series_equal(ser.nlargest(len(ser) + 1), ser.iloc[[4, 0, 1, 3, 2]])
 
+        tm.assert_series_equal(ser.nsorted(2, True), ser.nsmallest(2))
+        tm.assert_series_equal(ser.nsorted(2, False), ser.nlargest(2))
+
     def test_nlargest_misc(self):
         ser = Series([3.0, np.nan, 1, 2, 5])
         result = ser.nlargest()