How to write unittest for function that accepts csv file?

Viewed 1308

My function looks like this

import pandas as pd 

def my_function( file_path ):
    df = pd.read_csv( filepath )
    #some other operations on df here 
    return df 

I do not want to access an external file for unit test. Is there any way to create a csv file like object on the fly for unit testing.

1 Answers

Personally, I think it would be better to re-write your function so that it accepts a DataFrame, does some calculation on that DataFrame, and returns the DataFrame. In this sense you have made your function more atomic, it only manipulates the DataFrame as opposed to loading data and manipulating data. For this, you'd then create a test DataFrame and check the output. You can look at how pandas implements tests to get some inspiration

If you don't care for that, then we can make use of the flexibility of pd.read_csv and pass it a _io.StringIO object which it will parse into a DataFrame.

import pandas as pd
from io import StringIO

def my_function(file_path):
    df = pd.read_csv(filepath)
    return df 

file_path = StringIO('col1,col2,col3\n1,2,3\n4,5,6')

df = my_function(file_path)
print(df)
#   col1  col2  col3
#0     1     2     3
#1     4     5     6
Related