test changes

Author: Moi

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -287,12 +287,10 @@ def test_multi_function_flexible_mix(df):
  result = grouped.aggregate(d)
  tm.assert_frame_equal(result, expected)
 
-
+@pytest.mark.parametrize('as_index', [True, False])
 def test_not_as_index_agg_list():
 
  # GH 25011
- # expected behavhior of agg with as_index=True
- as_index = False
  array = [[3, 1, 2],
  [3, 3, 4],
  [4, 5, 6],
@@ -301,31 +299,23 @@ def test_not_as_index_agg_list():
  groupby = df.groupby('shouldnt_be_index', as_index=as_index)
  result = groupby.agg(['min', 'max'])
 
- array2 = [[3, 1, 3, 2, 4],
- [4, 5, 7, 6, 8]]
- columns = pd.MultiIndex(levels=[['A', 'B', 'shouldnt_be_index'],
- ['min', 'max', '']],
- codes=[[2, 0, 0, 1, 1], [2, 0, 1, 0, 1]])
- expected = pd.DataFrame(array2, columns=columns)
+ if as_index:
 
- tm.assert_frame_equal(result, expected)
+ array2 = [[1, 3, 2, 4],
+ [5, 7, 6, 8]]
+ columns = pd.MultiIndex(levels=[['A', 'B'],
+ ['min', 'max']],
+ codes=[[0, 0, 1, 1], [0, 1, 0, 1]])
+ index = pd.Series([3, 4], name='should_be_index')
+ expected = pd.DataFrame(array2, columns=columns, index=index)
 
- # expected behavhior of agg with as_index=True
- as_index = True
- array = [[3, 1, 2],
- [3, 3, 4],
- [4, 5, 6],
- [4, 7, 8]]
- df = pd.DataFrame(array, columns=['should_be_index', 'A', 'B'])
- groupby = df.groupby('should_be_index', as_index=as_index)
- result = groupby.agg(['min', 'max'])
+ else:
 
- array2 = [[1, 3, 2, 4],
- [5, 7, 6, 8]]
- columns = pd.MultiIndex(levels=[['A', 'B'],
- ['min', 'max']],
- codes=[[0, 0, 1, 1], [0, 1, 0, 1]])
- index = pd.Series([3, 4], name='should_be_index')
- expected = pd.DataFrame(array2, columns=columns, index=index)
+ array2 = [[3, 1, 3, 2, 4],
+ [4, 5, 7, 6, 8]]
+ columns = pd.MultiIndex(levels=[['A', 'B', 'shouldnt_be_index'],
+ ['min', 'max', '']],
+ codes=[[2, 0, 0, 1, 1], [2, 0, 1, 0, 1]])
+ expected = pd.DataFrame(array2, columns=columns)
 
  tm.assert_frame_equal(result, expected)