Add More Error Tests

pandas/tests/frame/test_analytics.py

@@ -1916,13 +1916,33 @@ def df_strings():
                          'c': np.random.permutation(10).astype('float64')})
 
 
+@pytest.fixture
+def df_main_dtypes():
+    return pd.DataFrame(
+        {'group': [1, 1, 2],
+         'int': [1, 2, 3],
+         'float': [4., 5., 6.],
+         'string': list('abc'),
+         'category_string': pd.Series(list('abc')).astype('category'),
+         'category_int': [7, 8, 9],
+         'datetime': pd.date_range('20130101', periods=3),
+         'datetimetz': pd.date_range('20130101',
+                                     periods=3,
+                                     tz='US/Eastern'),
+         'timedelta': pd.timedelta_range('1 s', periods=3, freq='s')},
+        columns=['group', 'int', 'float', 'string',
+                 'category_string', 'category_int',
+                 'datetime', 'datetimetz',
+                 'timedelta'])
+
+
 class TestNLargestNSmallest(object):
 
     # ----------------------------------------------------------------------
     # Top / bottom
     @pytest.mark.parametrize(
-        'n, order',
-        product(range(1, 11),
+        'method, n, order',
+        product(['nsmallest', 'nlargest'], range(1, 11),
                 [['a'],
                  ['c'],
                  ['a', 'b'],
@@ -1939,39 +1959,46 @@ class TestNLargestNSmallest(object):
                  ['b', 'c', 'c'],
 
                  ]))
-    def test_n(self, df_strings, n, order):
+    def test_n(self, df_strings, method, n, order):
         # GH10393
         df = df_strings
-
-        error_msg = (
-            "'b' has dtype: object, cannot use method 'nsmallest' "
-            "with this dtype"
-        )
-        if 'b' in order:
-            with pytest.raises(TypeError) as exception:
-                df.nsmallest(n, order)
-            assert exception.value, error_msg
+        if order[0] == 'b':
+
+            # Only expect error when 'b' is first in order, as 'a' and 'c' are
+            # unique
+            error_msg = (
+                "'b' has dtype: object, cannot use method 'nsmallest' "
+                "with this dtype"
+            )
+            with pytest.raises(TypeError) as exc_info:
+                getattr(df, method)(n, order)
+            assert exc_info.value, error_msg
         else:
-            result = df.nsmallest(n, order)
-            expected = df.sort_values(order).head(n)
+            ascending = method == 'nsmallest'
+            result = getattr(df, method)(n, order)
+            expected = df.sort_values(order, ascending=ascending).head(n)
             tm.assert_frame_equal(result, expected)
 
-        if 'b' in order:
-            with pytest.raises(TypeError) as exception:
-                df.nsmallest(n, order)
-            assert exception.value, error_msg
-        else:
-            result = df.nlargest(n, order)
-            expected = df.sort_values(order, ascending=False).head(n)
-            tm.assert_frame_equal(result, expected)
-
-    def test_n_error(self, df_strings):
-        # b alone raises a TypeError
-        df = df_strings
-        with pytest.raises(TypeError):
-            df.nsmallest(1, 'b')
-        with pytest.raises(TypeError):
-            df.nlargest(1, 'b')
+    @pytest.mark.parametrize(
+        'method, columns',
+        product(['nsmallest', 'nlargest'],
+                product(['group'], ['category_string', 'string'])
+                ))
+    def test_n_error(self, df_main_dtypes, method, columns):
+        df = df_main_dtypes
+        with pytest.raises(TypeError) as exc_info:
+            getattr(df, method)(2, columns)
+            msg = "Cannot use method '%s' with dtype %s" % (
+                method, df[columns[1]].dtype
+            )
+            assert exc_info.value, msg
+
+    def test_n_all_dtypes(self, df_main_dtypes):
+        df = df_main_dtypes
+        df.nsmallest(2, list(set(df) - {'category_string', 'string'}))
+        df.nsmallest(2, ['int', 'string'])  # int column is unique => OK
+        df.nlargest(2, list(set(df) - {'category_string', 'string'}))
+        df.nlargest(2, ['int', 'string'])  # int column is unique => OK
 
     def test_n_identical_values(self):
         # GH15297

pandas/tests/series/test_analytics.py

@@ -1667,7 +1667,9 @@ class TestNLargestNSmallest(object):
               Series([3., 2, 1, 2, 5], dtype='object'),
               # not supported on some archs
               # Series([3., 2, 1, 2, 5], dtype='complex256'),
-              Series([3., 2, 1, 2, 5], dtype='complex128')])
+              Series([3., 2, 1, 2, 5], dtype='complex128'),
+              Series(list('abcde'), dtype='category'),
+              Series(list('abcde'))])
     def test_error(self, r):
         dt = r.dtype
         msg = "Cannot use method 'n(larg|small)est' with dtype %s" % dt