Unit test for a function that transforms a PANDAS dataframe

Question

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_PERIODSTART'] += pd.to_timedelta(ops.groupby('TRADEID').cumcount(), unit='d')
        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()

Answer 1

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 = """
SEQ_PERIODSTART  SEQ_PERIODEND  TRADEID  NDAYS
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"""
expected = pd.read_csv(
    StringIO(expected_str),
    sep="\s+",
    parse_dates=["SEQ_PERIODSTART", "SEQ_PERIODEND"],
)

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.