Change all assert_frame/assert_series to tm.assert_

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -9,7 +9,6 @@
 import pandas as pd
 
 from pandas import concat, DataFrame, Index, MultiIndex, Series
-from pandas.util.testing import assert_frame_equal, assert_series_equal
 from pandas.core.groupby import SpecificationError
 from pandas.compat import OrderedDict
 import pandas.util.testing as tm
@@ -61,7 +60,7 @@ def test_agg_regression1(self):
         grouped = self.tsframe.groupby([lambda x: x.year, lambda x: x.month])
         result = grouped.agg(np.mean)
         expected = grouped.mean()
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_agg_must_agg(self):
         grouped = self.df.groupby('A')['C']
@@ -79,7 +78,7 @@ def test_agg_ser_multi_key(self):
         f = lambda x: x.sum()
         results = self.df.C.groupby([self.df.A, self.df.B]).aggregate(f)
         expected = self.df.groupby(['A', 'B']).sum()['C']
-        assert_series_equal(results, expected)
+        tm.assert_series_equal(results, expected)
 
     def test_agg_apply_corner(self):
         # nothing to group, all NA
@@ -90,18 +89,18 @@ def test_agg_apply_corner(self):
         exp = Series([],
                      dtype=np.float64,
                      index=pd.Index([], dtype=np.float64))
-        assert_series_equal(grouped.sum(), exp)
-        assert_series_equal(grouped.agg(np.sum), exp)
-        assert_series_equal(grouped.apply(np.sum), exp, check_index_type=False)
+        tm.assert_series_equal(grouped.sum(), exp)
+        tm.assert_series_equal(grouped.agg(np.sum), exp)
+        tm.assert_series_equal(grouped.apply(np.sum), exp, check_index_type=False)
 
         # DataFrame
         grouped = self.tsframe.groupby(self.tsframe['A'] * np.nan)
         exp_df = DataFrame(columns=self.tsframe.columns,
                            dtype=float,
                            index=pd.Index([], dtype=np.float64))
-        assert_frame_equal(grouped.sum(), exp_df, check_names=False)
-        assert_frame_equal(grouped.agg(np.sum), exp_df, check_names=False)
-        assert_frame_equal(grouped.apply(np.sum), exp_df.iloc[:, :0],
+        tm.assert_frame_equal(grouped.sum(), exp_df, check_names=False)
+        tm.assert_frame_equal(grouped.agg(np.sum), exp_df, check_names=False)
+        tm.assert_frame_equal(grouped.apply(np.sum), exp_df.iloc[:, :0],
                            check_names=False)
 
     def test_agg_grouping_is_list_tuple(self):
@@ -135,20 +134,20 @@ def _check_results(grouped):
             # single series
             result = grouped['A'].agg('std')
             expected = grouped['A'].std()
-            assert_series_equal(result, expected)
+            tm.assert_series_equal(result, expected)
 
             # group frame by function name
             result = grouped.aggregate('var')
             expected = grouped.var()
-            assert_frame_equal(result, expected)
+            tm.assert_frame_equal(result, expected)
 
             # group frame by function dict
             result = grouped.agg(OrderedDict([['A', 'var'], ['B', 'std'],
                                               ['C', 'mean'], ['D', 'sem']]))
             expected = DataFrame(OrderedDict([['A', grouped['A'].var(
             )], ['B', grouped['B'].std()], ['C', grouped['C'].mean()],
                 ['D', grouped['D'].sem()]]))
-            assert_frame_equal(result, expected)
+            tm.assert_frame_equal(result, expected)
 
         by_weekday = self.tsframe.groupby(lambda x: x.weekday())
         _check_results(by_weekday)
@@ -204,7 +203,7 @@ def func(ser):
         result = grouped.aggregate(func)
         exp_grouped = self.three_group.loc[:, self.three_group.columns != 'C']
         expected = exp_grouped.groupby(['A', 'B']).aggregate(func)
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_agg_multiple_functions_maintain_order(self):
         # GH #610
@@ -222,11 +221,11 @@ def test_multiple_functions_tuples_and_non_tuples(self):
 
         result = self.df.groupby('A')['C'].agg(funcs)
         expected = self.df.groupby('A')['C'].agg(ex_funcs)
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
         result = self.df.groupby('A').agg(funcs)
         expected = self.df.groupby('A').agg(ex_funcs)
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_agg_multiple_functions_too_many_lambdas(self):
         grouped = self.df.groupby('A')
@@ -249,14 +248,14 @@ def test_more_flexible_frame_multi_function(self):
         d = OrderedDict([['C', [np.mean, np.std]], ['D', [np.mean, np.std]]])
         result = grouped.aggregate(d)
 
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
         # be careful
         result = grouped.aggregate(OrderedDict([['C', np.mean],
                                                 ['D', [np.mean, np.std]]]))
         expected = grouped.aggregate(OrderedDict([['C', np.mean],
                                                   ['D', [np.mean, np.std]]]))
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
         def foo(x):
             return np.mean(x)
@@ -274,7 +273,7 @@ def bar(x):
         d = OrderedDict([['C', [np.mean]], ['D', [foo, bar]]])
         expected = grouped.aggregate(d)
 
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_multi_function_flexible_mix(self):
         # GH #1268
@@ -306,5 +305,5 @@ def test_multi_function_flexible_mix(self):
                                         check_stacklevel=False):
             expected = grouped.aggregate(d3)
 
-        assert_frame_equal(result, expected)
-        assert_frame_equal(result2, expected)
+        tm.assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result2, expected)

pandas/tests/groupby/aggregate/test_cython.py

@@ -12,7 +12,6 @@
 import pandas as pd
 
 from pandas import bdate_range, DataFrame, Index, Series
-from pandas.util.testing import assert_frame_equal, assert_series_equal
 from pandas.core.groupby import DataError
 import pandas.util.testing as tm
 
@@ -38,7 +37,7 @@ def _testit(name):
             exp = DataFrame({'C': exp})
             exp.index.name = 'A'
             result = op(grouped)
-            assert_frame_equal(result, exp)
+            tm.assert_frame_equal(result, exp)
 
             # multiple columns
             grouped = df.groupby(['A', 'B'])
@@ -51,7 +50,7 @@ def _testit(name):
 
             result = op(grouped)['C']
             if name in ['sum', 'prod']:
-                assert_series_equal(result, exp)
+                tm.assert_series_equal(result, exp)
 
         _testit('count')
         _testit('sum')
@@ -70,7 +69,7 @@ def test_cython_agg_boolean(self):
         result = frame.groupby('a')['b'].mean()
         expected = frame.groupby('a')['b'].agg(np.mean)
 
-        assert_series_equal(result, expected)
+        tm.assert_series_equal(result, expected)
 
     def test_cython_agg_nothing_to_agg(self):
         frame = DataFrame({'a': np.random.randint(0, 5, 50),
@@ -117,7 +116,7 @@ def test_cython_agg_return_dict(self):
                            {'two': 2, 'one': 2, 'three': 1}],
                           index=Index(['bar', 'foo'], name='A'),
                           name='B')
-        assert_series_equal(ts, expected)
+        tm.assert_series_equal(ts, expected)
 
     def test_cython_fail_agg(self):
         dr = bdate_range('1/1/2000', periods=50)
@@ -126,7 +125,7 @@ def test_cython_fail_agg(self):
         grouped = ts.groupby(lambda x: x.month)
         summed = grouped.sum()
         expected = grouped.agg(np.sum)
-        assert_series_equal(summed, expected)
+        tm.assert_series_equal(summed, expected)
 
     def test__cython_agg_general(self):
         ops = [('mean', np.mean),

pandas/tests/groupby/aggregate/test_other.py

@@ -16,7 +16,6 @@
 import pandas as pd
 
 from pandas import date_range, DataFrame, Index, MultiIndex, Series
-from pandas.util.testing import assert_frame_equal, assert_series_equal
 from pandas.core.groupby import SpecificationError
 from pandas.io.formats.printing import pprint_thing
 import pandas.util.testing as tm
@@ -39,7 +38,7 @@ def peak_to_peak(arr):
     expected = grouped.agg([peak_to_peak])
     expected.columns = ['data1', 'data2']
     result = grouped.agg(peak_to_peak)
-    assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_agg_datetimes_mixed():
@@ -93,32 +92,32 @@ def test_agg_dict_parameter_cast_result_dtypes():
     # test for `first` function
     exp = df.loc[[0, 3, 4, 6]].set_index('class')
     grouped = df.groupby('class')
-    assert_frame_equal(grouped.first(), exp)
-    assert_frame_equal(grouped.agg('first'), exp)
-    assert_frame_equal(grouped.agg({'time': 'first'}), exp)
-    assert_series_equal(grouped.time.first(), exp['time'])
-    assert_series_equal(grouped.time.agg('first'), exp['time'])
+    tm.assert_frame_equal(grouped.first(), exp)
+    tm.assert_frame_equal(grouped.agg('first'), exp)
+    tm.assert_frame_equal(grouped.agg({'time': 'first'}), exp)
+    tm.assert_series_equal(grouped.time.first(), exp['time'])
+    tm.assert_series_equal(grouped.time.agg('first'), exp['time'])
 
     # test for `last` function
     exp = df.loc[[0, 3, 4, 7]].set_index('class')
     grouped = df.groupby('class')
-    assert_frame_equal(grouped.last(), exp)
-    assert_frame_equal(grouped.agg('last'), exp)
-    assert_frame_equal(grouped.agg({'time': 'last'}), exp)
-    assert_series_equal(grouped.time.last(), exp['time'])
-    assert_series_equal(grouped.time.agg('last'), exp['time'])
+    tm.assert_frame_equal(grouped.last(), exp)
+    tm.assert_frame_equal(grouped.agg('last'), exp)
+    tm.assert_frame_equal(grouped.agg({'time': 'last'}), exp)
+    tm.assert_series_equal(grouped.time.last(), exp['time'])
+    tm.assert_series_equal(grouped.time.agg('last'), exp['time'])
 
     # count
     exp = pd.Series([2, 2, 2, 2],
                     index=Index(list('ABCD'), name='class'),
                     name='time')
-    assert_series_equal(grouped.time.agg(len), exp)
-    assert_series_equal(grouped.time.size(), exp)
+    tm.assert_series_equal(grouped.time.agg(len), exp)
+    tm.assert_series_equal(grouped.time.size(), exp)
 
     exp = pd.Series([0, 1, 1, 2],
                     index=Index(list('ABCD'), name='class'),
                     name='time')
-    assert_series_equal(grouped.time.count(), exp)
+    tm.assert_series_equal(grouped.time.count(), exp)
 
 
 def test_agg_cast_results_dtypes():
@@ -130,7 +129,7 @@ def test_agg_cast_results_dtypes():
 
     result = df.groupby('X')['Y'].agg(len)
     expected = df.groupby('X')['Y'].count()
-    assert_series_equal(result, expected)
+    tm.assert_series_equal(result, expected)
 
 
 def test_aggregate_float64_no_int64():
@@ -177,7 +176,7 @@ def test_aggregate_api_consistency():
     expected = pd.concat([d_sum, d_mean],
                          axis=1)
     expected.columns = ['sum', 'mean']
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     result = grouped.agg([np.sum, np.mean])
     expected = pd.concat([c_sum,
@@ -187,7 +186,7 @@ def test_aggregate_api_consistency():
                          axis=1)
     expected.columns = MultiIndex.from_product([['C', 'D'],
                                                 ['sum', 'mean']])
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     result = grouped[['D', 'C']].agg([np.sum, np.mean])
     expected = pd.concat([d_sum,
@@ -197,13 +196,13 @@ def test_aggregate_api_consistency():
                          axis=1)
     expected.columns = MultiIndex.from_product([['D', 'C'],
                                                 ['sum', 'mean']])
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     result = grouped.agg({'C': 'mean', 'D': 'sum'})
     expected = pd.concat([d_sum,
                           c_mean],
                          axis=1)
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     result = grouped.agg({'C': ['mean', 'sum'],
                           'D': ['mean', 'sum']})
@@ -226,7 +225,7 @@ def test_aggregate_api_consistency():
                          axis=1)
     expected.columns = MultiIndex.from_product([['r', 'r2'],
                                                 ['D', 'C']])
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
 
 def test_agg_dict_renaming_deprecation():
@@ -270,7 +269,7 @@ def test_agg_compat():
     with tm.assert_produces_warning(FutureWarning,
                                     check_stacklevel=False):
         result = g['D'].agg({'C': ['sum', 'std']})
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     expected = pd.concat([g['D'].sum(),
                           g['D'].std()],
@@ -280,7 +279,7 @@ def test_agg_compat():
     with tm.assert_produces_warning(FutureWarning,
                                     check_stacklevel=False):
         result = g['D'].agg({'C': 'sum', 'D': 'std'})
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
 
 def test_agg_nested_dicts():
@@ -310,7 +309,7 @@ def test_agg_nested_dicts():
     expected.columns = pd.MultiIndex.from_tuples(
         [('ra', 'mean'), ('ra', 'std'),
          ('rb', 'mean'), ('rb', 'std')])
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
     # same name as the original column
     # GH9052
@@ -322,7 +321,7 @@ def test_agg_nested_dicts():
     with tm.assert_produces_warning(FutureWarning,
                                     check_stacklevel=False):
         result = g['D'].agg({'D': np.sum, 'result2': np.mean})
-    assert_frame_equal(result, expected, check_like=True)
+    tm.assert_frame_equal(result, expected, check_like=True)
 
 
 def test_agg_item_by_item_raise_typeerror():
@@ -345,7 +344,7 @@ def test_series_agg_multikey():
 
     result = grouped.agg(np.sum)
     expected = grouped.sum()
-    assert_series_equal(result, expected)
+    tm.assert_series_equal(result, expected)
 
 
 def test_series_agg_multi_pure_python():
@@ -366,7 +365,7 @@ def bad(x):
 
     result = data.groupby(['A', 'B']).agg(bad)
     expected = data.groupby(['A', 'B']).agg(lambda x: 'foo')
-    assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_agg_consistency():
@@ -391,7 +390,7 @@ def P1(a):
     expected.columns = expected.columns.levels[0]
 
     result = g.agg(P1)
-    assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_agg_callables():
@@ -409,7 +408,7 @@ def __call__(self, x):
     expected = df.groupby("foo").agg(sum)
     for ecall in equiv_callables:
         result = df.groupby('foo').agg(ecall)
-        assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
 
 def test_agg_over_numpy_arrays():
@@ -427,7 +426,7 @@ def test_agg_over_numpy_arrays():
                             index=expected_index,
                             columns=expected_column)
 
-    assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_agg_timezone_round_trip():
@@ -493,7 +492,7 @@ def test_agg_structs_dataframe(structure, expected):
 
     result = df.groupby(['A', 'B']).aggregate(structure)
     expected.index.names = ['A', 'B']
-    assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 @pytest.mark.parametrize("structure, expected", [
@@ -512,4 +511,4 @@ def test_agg_structs_series(structure, expected):
 
     result = df.groupby('A')['C'].aggregate(structure)
     expected.index.name = 'A'
-    assert_series_equal(result, expected)
+    tm.assert_series_equal(result, expected)