# Unit test for a function that transforms a PANDAS dataframe

I am currently working on a project that uses many dataframe wrangling. The intention is to write test to check that the output functions are correct.

I have written a test for a dataframe function that transforms a dataframe of operations into another dataframe that contains the volumen of the operations for each of the containing days.

My question is:

### Is this a good test for a function that receives a DataFrame and returns a DataFrame?

Should I do something different?

This is is my code:

import pandas as pd

class Operations:

def disaggregate_ops_volume_date(self, df: DataFrame):
""" Given a DataFrame with operations it generates an operation for each of the days it contains """

if df.empty:
return df

ops = df.copy()

ops['NDAYS'] = ops[['SEQ_PERIODSTART', 'SEQ_PERIODEND']].apply(
lambda row: (row['SEQ_PERIODEND'] - row['SEQ_PERIODSTART']).days, axis=1)

# Add a copy of each operation NDAYS times
ops = ops.loc[np.repeat(ops.index, ops['NDAYS'])]

# Correct the date based on the TradeId operation
ops['SEQ_PERIODEND'] = ops['SEQ_PERIODSTART'] + pd.Timedelta(1, unit='d')

ops = ops.reset_index(drop=True)

return ops


### This is the test I have developed:

from pandas.util.testing import assert_frame_equal
import unittest
import pandas as pd

class TestOperations(unittest.TestCase):

def test_minimal(self):
"""To make sure at least a test is passed"""
self.assertEqual('foo'.upper(), 'FOO')

def test_disaggregate_ops_volume_date(self):
input = pd.DataFrame(
{'SEQ_PERIODSTART':
['2019-02-10', '2019-02-12', '2019-02-13'],
'SEQ_PERIODEND':
['2019-02-11', '2019-02-14', '2019-02-18'],
'ID': [0, 1, 2]})

input['SEQ_PERIODSTART'] = pd.to_datetime(input['SEQ_PERIODSTART'])
input['SEQ_PERIODEND'] = pd.to_datetime(input['SEQ_PERIODEND'])

expected = pd.DataFrame(
{'SEQ_PERIODSTART':
['2019-02-10', '2019-02-12',                                                           '2019-02-13', '2019-02-13',                                              '2019-02-14', '2019-02-15',                                             '2019-02-16', '2019-02-17'],

'SEQ_PERIODEND':
['2019-02-11',  '2019-02-13',
'2019-02-14', '2019-02-14',                                               '2019-02-15', '2019-02-16',                                          '2019-02-17', '2019-02-18'],
'ID': [0, 1, 2, 3, 4, 5, 6, 7],
'NDAYS': [1, 2, 2, 5, 5, 5, 5, 5]
})
expected['SEQ_PERIODSTART'] = pd.to_datetime(expected['SEQ_PERIODSTART'])
expected['SEQ_PERIODEND'] = pd.to_datetime(expected['SEQ_PERIODEND'])

myops = Operations()

assert_frame_equal(expected, myops.disaggregate_ops_volume_date(input))

if __name__ == '__main__':
unittest.main()


The test itself seems okay.

# bug

You use "ID" as columns name in the test, but "TRADEID" in the method to test.

And the result is not what was expected. The TRADEIDs don't match

# Classes

There is no use for the class Operations. If you want to group functions that belong together, you can do so in a module (file) and import that. The fact disaggregate_ops_volume_date has a self parameter that is unused is a giveaway here

# Vectorize

there is no reason to calculate ops['NDAYS'] row per row via apply. (df["SEQ_PERIODEND"] - df["SEQ_PERIODSTART"]).dt.days works just as well.

Then you can use DataFrame.assign, and don't have to explicitly make a copy of df

ndays = (df["SEQ_PERIODEND"] - df["SEQ_PERIODSTART"]).dt.days
ops = df.assign(NDAYS = ndays)


# Day

Since a days is used a lot in that method, it can be clearer to define that up front:

DAY = pd.Timedelta("1d")


ndays can then be defined as ndays = (df["SEQ_PERIODEND"] - df["SEQ_PERIODSTART"]) // DAY. Whether this is clearer than the .dt.days is a matter of taste.

correcting the date can then be expessed more clearly as:

ops["SEQ_PERIODSTART"] += ops.groupby("TRADEID").cumcount() * DAY
ops["SEQ_PERIODEND"] = ops["SEQ_PERIODSTART"] + DAY


# input

input is a builtin. By using that name as a variable name, you shadow that builtin. In this case this is not a big problem, but in general you should avoid this

# indentation style

You use a very inconsistent style of indenting the code, and what goes on a separate line. Better would be to remain consistent. For this, I use black.

This changes:

expected = pd.DataFrame(
{'SEQ_PERIODSTART':
['2019-02-10', '2019-02-12',                                                           '2019-02-13', '2019-02-13',                                              '2019-02-14', '2019-02-15',                                             '2019-02-16', '2019-02-17'],

'SEQ_PERIODEND':
['2019-02-11',  '2019-02-13',
'2019-02-14', '2019-02-14',                                               '2019-02-15', '2019-02-16',                                          '2019-02-17', '2019-02-18'],
'ID': [0, 1, 2, 3, 4, 5, 6, 7],
'NDAYS': [1, 2, 2, 5, 5, 5, 5, 5]
})


to:

expected = pd.DataFrame(
{
"SEQ_PERIODSTART": [
"2019-02-10",
"2019-02-12",
"2019-02-13",
"2019-02-13",
"2019-02-14",
"2019-02-15",
"2019-02-16",
"2019-02-17",
],
"SEQ_PERIODEND": [
"2019-02-11",
"2019-02-13",
"2019-02-14",
"2019-02-14",
"2019-02-15",
"2019-02-16",
"2019-02-17",
"2019-02-18",
],
"ID": [0, 1, 2, 3, 4, 5, 6, 7],
"NDAYS": [1, 2, 2, 5, 5, 5, 5, 5],
}
)


Which I think, is a lot more clear

# pd.to_datetime

You can invoke this immediately on a list, so defining the input DataFrame can become:

input_df = pd.DataFrame(
{
"SEQ_PERIODSTART": pd.to_datetime(
["2019-02-10", "2019-02-12", "2019-02-13"]
),
"SEQ_PERIODEND": pd.to_datetime(
["2019-02-11", "2019-02-14", "2019-02-18"]
),
"ID": [0, 1, 2],
}
)


# defining test DataFrames

Instead of invoking pd.DataFrame immediately, an alternative is working via a csv-like text input, and then use pd.read_csv and StringIO to convert it to a DataFrame.

expected_str = """
2019-02-10       2019-02-11     0        1
2019-02-12       2019-02-13     1        2
2019-02-13       2019-02-14     2        2
2019-02-13       2019-02-14     3        5
2019-02-14       2019-02-15     4        5
2019-02-15       2019-02-16     5        5
2019-02-16       2019-02-17     6        5
2019-02-17       2019-02-18     7        5"""

Whether this is more clear than using pd.DataFrame() is a matter of taste. The advantage of this method is that it is easy to add or remove a line, and see whether the data is aligned correctly.